core.models.uma.nn.layer_norm

Contents

core.models.uma.nn.layer_norm#

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.

1. Normalize features of shape (N, sphere_basis, C), with sphere_basis = (lmax + 1) ** 2.

2. The difference from layer_norm.py is that all type-L vectors have the same number of channels and input features are of shape (N, sphere_basis, C).

Classes#

EquivariantLayerNormArray

Base class for all neural network modules.

EquivariantLayerNormArraySphericalHarmonics

EquivariantRMSNormArraySphericalHarmonics

EquivariantRMSNormArraySphericalHarmonicsV2

EquivariantDegreeLayerScale

Functions#

get_normalization_layer(norm_type, lmax, num_channels)

get_l_to_all_m_expand_index(lmax)

Module Contents#

core.models.uma.nn.layer_norm.get_normalization_layer(norm_type: str, lmax: int, num_channels: int, eps: float = 1e-05, affine: bool = True, normalization: str = 'component')#
core.models.uma.nn.layer_norm.get_l_to_all_m_expand_index(lmax: int)#
class core.models.uma.nn.layer_norm.EquivariantLayerNormArray(lmax: int, num_channels: int, eps: float = 1e-05, affine: bool = True, normalization: str = 'component')#

Bases: torch.nn.Module

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:

training (bool) – Boolean represents whether this module is in training or evaluation mode.

lmax#
num_channels#
eps#
affine#
normalization#
__repr__() → str#
forward(node_input)#

Assume input is of shape [N, sphere_basis, C]

class core.models.uma.nn.layer_norm.EquivariantLayerNormArraySphericalHarmonics(lmax: int, num_channels: int, eps: float = 1e-05, affine: bool = True, normalization: str = 'component', std_balance_degrees: bool = True)#

Bases: torch.nn.Module

  1. Normalize over L = 0.

  2. Normalize across all m components from degrees L > 0.

  3. Do not normalize separately for different L (L > 0).

lmax#
num_channels#
eps#
affine#
std_balance_degrees#
norm_l0#
normalization#
__repr__() → str#
forward(node_input)#

Assume input is of shape [N, sphere_basis, C]

class core.models.uma.nn.layer_norm.EquivariantRMSNormArraySphericalHarmonics(lmax: int, num_channels: int, eps: float = 1e-05, affine: bool = True, normalization: str = 'component')#

Bases: torch.nn.Module

  1. Normalize across all m components from degrees L >= 0.

lmax#
num_channels#
eps#
affine#
normalization#
__repr__() → str#
forward(node_input)#

Assume input is of shape [N, sphere_basis, C]

class core.models.uma.nn.layer_norm.EquivariantRMSNormArraySphericalHarmonicsV2(lmax: int, num_channels: int, eps: float = 1e-05, affine: bool = True, normalization: str = 'component', centering: bool = True, std_balance_degrees: bool = True)#

Bases: torch.nn.Module

  1. Normalize across all m components from degrees L >= 0.

  2. Expand weights and multiply with normalized feature to prevent slicing and concatenation.

lmax#
num_channels#
eps#
affine#
centering#
std_balance_degrees#
normalization#
__repr__() → str#
forward(node_input)#

Assume input is of shape [N, sphere_basis, C]

class core.models.uma.nn.layer_norm.EquivariantDegreeLayerScale(lmax: int, num_channels: int, scale_factor: float = 2.0)#

Bases: torch.nn.Module

  1. Similar to Layer Scale used in CaiT (Going Deeper With Image Transformers (ICCV’21)), we scale the output of both attention and FFN.

  2. For degree L > 0, we scale down the square root of 2 * L, which is to emulate halving the number of channels when using higher L.

lmax#
num_channels#
scale_factor#
affine_weight#
__repr__() → str#
forward(node_input)#
Contents