xnn.gnn.models.blocks.RadialNet#

class xnn.gnn.models.blocks.RadialNet(hs, act=None, variance_in=1, variance_out=1, out_act=False)[source]#

Bases: FullyConnectedNet

Per-edge radial network that can recompute its hidden layers in the backward.

An e3nn.nn.FullyConnectedNet (same parameters, state_dict keys and outputs). When a gradient flows through it (forces, stress), the hidden layers’ activations, E x hidden per layer whatever the tensor product’s size, are kept for the backward pass; with recomputation only the last hidden output is kept and the hidden layers run again in the backward, giving the same values. TorchScript always keeps the activations.

Variables:: recompute (str) – "auto" (default): recompute in evaluation mode from RECOMPUTE_MIN_EDGES edges; "on": whenever a gradient is needed; "off": never. Set with set_recompute_radial().
Parameters:: hs (List[int])

forward(x)[source]#

Apply the network to x of shape (E, hs[0]).

Parameters:: x (Tensor)
Return type:: Tensor