core.models.equiformer_v2.so2_ops

Classes

SO2_m_Convolution	SO(2) Conv: Perform an SO(2) convolution on features corresponding to +- m
SO2_Convolution	SO(2) Block: Perform SO(2) convolutions for all m (orders)
SO2_Linear	SO(2) Linear: Perform SO(2) linear for all m (orders).

Module Contents

class core.models.equiformer_v2.so2_ops.SO2_m_Convolution(m: int, sphere_channels: int, m_output_channels: int, lmax_list: list[int], mmax_list: list[int])

Bases: torch.nn.Module

SO(2) Conv: Perform an SO(2) convolution on features corresponding to +- m

Parameters:

• m (int) – Order of the spherical harmonic coefficients

• sphere_channels (int) – Number of spherical channels

• m_output_channels (int) – Number of output channels used during the SO(2) conv

• (list (mmax_list) – int): List of degrees (l) for each resolution

• (list – int): List of orders (m) for each resolution

m

sphere_channels

m_output_channels

lmax_list

mmax_list

num_resolutions: int

fc

forward(x_m)

class core.models.equiformer_v2.so2_ops.SO2_Convolution(sphere_channels: int, m_output_channels: int, lmax_list: list[int], mmax_list: list[int], mappingReduced, internal_weights: bool = True, edge_channels_list: list[int] | None = None, extra_m0_output_channels: int | None = None)

Bases: torch.nn.Module

SO(2) Block: Perform SO(2) convolutions for all m (orders)

Parameters:

• sphere_channels (int) – Number of spherical channels

• m_output_channels (int) – Number of output channels used during the SO(2) conv

• (list (edge_channels_list) – int): List of degrees (l) for each resolution

• (list – int): List of orders (m) for each resolution

• mappingReduced (CoefficientMappingModule) – Used to extract a subset of m components

• internal_weights (bool) – If True, not using radial function to multiply inputs features

• (list – int): List of sizes of invariant edge embedding. For example, [input_channels, hidden_channels, hidden_channels].

• extra_m0_output_channels (int) – If not None, return out_embedding (SO3_Embedding) and extra_m0_features (Tensor).

sphere_channels

m_output_channels

lmax_list

mmax_list

mappingReduced

num_resolutions

internal_weights

edge_channels_list

extra_m0_output_channels

fc_m0

so2_m_conv

rad_func = None

forward(x, x_edge)

class core.models.equiformer_v2.so2_ops.SO2_Linear(sphere_channels: int, m_output_channels: int, lmax_list: list[int], mmax_list: list[int], mappingReduced, internal_weights: bool = False, edge_channels_list: list[int] | None = None)

Bases: torch.nn.Module

SO(2) Linear: Perform SO(2) linear for all m (orders).

Parameters:

• sphere_channels (int) – Number of spherical channels

• m_output_channels (int) – Number of output channels used during the SO(2) conv

• (list (edge_channels_list) – int): List of degrees (l) for each resolution

• (list – int): List of orders (m) for each resolution

• mappingReduced (CoefficientMappingModule) – Used to extract a subset of m components

• internal_weights (bool) – If True, not using radial function to multiply inputs features

• (list – int): List of sizes of invariant edge embedding. For example, [input_channels, hidden_channels, hidden_channels].

sphere_channels

m_output_channels

lmax_list

mmax_list

mappingReduced

internal_weights

edge_channels_list

num_resolutions

fc_m0

so2_m_fc

rad_func = None

forward(x, x_edge)