BatchNorm#

class cuequivariance_torch.layers.BatchNorm#

Batch normalization for orthonormal representations.

It normalizes by the norm of the representations. Note that the norm is invariant only for orthonormal representations.

Parameters:

  • irreps (Irreps) – Input irreps.

  • layout (IrrepsLayout, optional) – Layout of the input tensor, by default IrrepsLayout.mul_ir.

  • eps (float, optional) – Epsilon value for numerical stability, by default 1e-5.

  • momentum (float, optional) – Momentum for the running mean and variance, by default 0.1.

  • affine (bool, optional) – Whether to apply an affine transformation, by default True.

  • reduce (str, optional) – How to reduce the norm of the representations, by default “mean”.

  • instance (bool, optional) – Whether to use instance normalization, by default False.

  • include_bias (bool, optional) – Whether to include a bias term, by default True.

Forward Pass

forward(input)#

Normalize the input tensor.

Parameters:

input (torch.Tensor) – Input tensor. The last dimension should match with the input irreps.

Returns:

Normalized tensor.

Return type:

torch.Tensor

__init__(
irreps,
*,
layout=None,
eps=1e-05,
momentum=0.1,
affine=True,
reduce='mean',
instance=False,
include_bias=True,
)#

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:
classmethod __new__(*args, **kwargs)#