mindscience.e3nn.nn.NormActivation

class mindscience.e3nn.nn.NormActivation(irreps_in, act, normalize=True, epsilon=None, bias=False, init_method='zeros', dtype=mindspore.float32, ncon_dtype=mindspore.float32)[source]

Activation function for the norm of irreps. Applies a scalar activation to the norm of each irrep and outputs a (normalized) version of that irrep multiplied by the scalar output of the scalar activation. Optionally, a learnable bias can be added to the norms before the activation, and the resulting features can be normalized by their original norm to preserve angular information while only modulating their magnitude.

Parameters

  • irreps_in (Union[str, Irrep, Irreps]) – Input irreps.

  • act (Func) – Activation function applied to the norm of each irrep.

  • normalize (bool, optional) – Whether to normalize input features before multiplying by the scalars from the nonlinearity. Default: True.

  • epsilon (float, optional) – When normalize, norms smaller than epsilon are clamped to epsilon to prevent division by zero. Ignored if normalize is False. Default: None.

  • bias (bool, optional) – Whether to apply a learnable additive bias to the inputs of act. Default: False.

  • init_method (Union[str, float, mindspore.common.initializer], optional) – Parameter initialization method. Default: 'zeros'.

  • dtype (mindspore.dtype, optional) – Data type of input tensors. Default: mindspore.float32.

  • ncon_dtype (mindspore.dtype, optional) – Data type for ncon computation. Default: mindspore.float32.

Inputs:

  • input (Tensor) - The shape of Tensor is \((..., irreps\_in.dim)\).

Outputs:

  • output (Tensor) - The shape of Tensor is \((..., irreps\_in.dim)\).

Raises

  • ValueError – If epsilon is not None and normalize is False.

  • ValueError – If epsilon is not positive.

Examples

>>> from mindscience.e3nn.nn import NormActivation
>>> from mindspore import ops, Tensor
>>> norm_activation = NormActivation("2x1e", ops.sigmoid, bias=True)
>>> print(norm_activation)
NormActivation [sigmoid] (2x1e -> 2x1e)
>>> inputs = Tensor(ops.ones((4, 6)))
>>> outputs = norm_activation(inputs)
>>> print(outputs.shape)
(4, 6)