xnn.gnn.models.blocks.ScalarActivation#

class xnn.gnn.models.blocks.ScalarActivation(irreps_in, act)[source]#

Bases: Module

Scriptable stand-in for e3nn.nn.Activation on all-scalar irreps.

e3nn.nn.Activation (0.4.4) does not compile under torch.jit.script on torch 2.x, but for all-scalar irreps – the only case the readouts and gates in this package need – it reduces to applying the second-moment- normalized activation (e3nn.math.normalize2mom()) elementwise. This module does exactly that, so it is numerically identical to the e3nn original while remaining TorchScript-compatible. It carries no state, so swapping it in leaves the state_dict layout untouched.

Parameters:

  • irreps_in (e3nn.o3.Irreps) – Irreps of the activated features; every entry must be l = 0.

  • act (callable or None) – Scalar activation; None means identity.

Raises:

ValueError – If irreps_in contains any l > 0 irrep.

forward(x)[source]#

Apply the normalized scalar activation (identity when act is None).

Parameters:

x (torch.Tensor) – All-scalar features.

Returns:

The activated features.

Return type:

torch.Tensor