xnn.gnn.models.blocks

Shared E(3)-equivariant building blocks (e3nn) for the GNN potentials.

Holds the pieces used by more than one equivariant model (NequIP / MACE / Allegro): irreps helpers, the tensor-product path test, a TorchScript-safe normalized scalar activation and the per-edge radial network. Model-specific blocks (MACE’s interaction / product blocks, NequIP’s convnet layer, …) live in the model modules. Requires e3nn.

Module Attributes

RECOMPUTE_MIN_EDGES

Rows (edges) from which a radial network in evaluation mode recomputes its hidden layers in the backward pass instead of keeping their activations (recompute="auto").

Functions

hidden_irreps(mul, l_max)	Hidden feature irreps: mul copies of the spherical-harmonic irreps.
set_recompute_radial(model[, recompute])	Choose when the radial networks of model recompute their hidden layers.
shifted_softplus(x)	Shifted softplus log(0.5 e^x + 0.5) (NequIP's ShiftedSoftPlus).
species_irreps(n_species)	Irreps of the one-hot species (node attribute) channels.
tp_out_irreps_with_instructions(irreps1, ...)	(uvu) tensor-product output irreps + instructions, keeping only target paths.
tp_path_exists(irreps_in1, irreps_in2, ir_out)	Whether some tensor-product path irreps_in1 x irreps_in2 -> ir_out exists.

Classes

RadialNet(hs[, act, variance_in, ...])	Per-edge radial network that can recompute its hidden layers in the backward.
ScalarActivation(irreps_in, act)	Scriptable stand-in for e3nn.nn.Activation on all-scalar irreps.