core.models.uma.nn.so2_layers

core.models.uma.nn.so2_layers#

Copyright (c) Meta Platforms, Inc. and affiliates.

This source code is licensed under the MIT license found in the LICENSE file in the root directory of this source tree.

Classes#

SO2_m_Conv

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.uma.nn.so2_layers.SO2_m_Conv(m: int, sphere_channels: int, m_output_channels: int, lmax: int, mmax: 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

  • lmax (int) – degrees (l)

  • mmax (int) – orders (m)

m#
sphere_channels#
m_output_channels#
lmax#
mmax#
out_channels_half#
fc#
forward(x_m)#
class core.models.uma.nn.so2_layers.SO2_Convolution(sphere_channels: int, m_output_channels: int, lmax: int, mmax: 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

  • lmax (int) – degrees (l)

  • mmax (int) – orders (m)

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

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

  • (list (edge_channels_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 and extra_m0_features (Tensor).

sphere_channels#
m_output_channels#
lmax#
mmax#
mappingReduced#
internal_weights#
extra_m0_output_channels#
fc_m0#
so2_m_conv#
rad_func = None#
m_split_sizes#
edge_split_sizes#
forward(x: torch.Tensor, x_edge: torch.Tensor)#
class core.models.uma.nn.so2_layers.SO2_Linear(sphere_channels: int, m_output_channels: int, lmax: int, mmax: 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

  • lmax (int) – degrees (l)

  • mmax (int) – orders (m)

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

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

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

sphere_channels#
m_output_channels#
lmax#
mmax#
mappingReduced#
internal_weights#
edge_channels_list#
fc_m0#
so2_m_fc#
rad_func = None#
forward(x, x_edge)#
Contents