xnn.gnn.models.painn.GatedEquivariantBlock#

class xnn.gnn.models.painn.GatedEquivariantBlock(n_in_s, n_in_v, n_out_s, n_out_v, n_hidden=None, scalar_activation=True)[source]#

Bases: Module

Gated equivariant block (paper Fig. 3).

Two linear maps of the vector features give W_1 v and W_2 v; the norm of the first is stacked with the scalars and fed through a two-layer network whose output is split into new scalars and a gate that scales W_2 v. Scalars pass through a nonlinearity, vectors are only scaled, so the block is equivariant.

Parameters:

  • n_in_s (int) – Input scalar and vector widths.

  • n_in_v (int) – Input scalar and vector widths.

  • n_out_s (int) – Output scalar and vector widths.

  • n_out_v (int) – Output scalar and vector widths.

  • n_hidden (int or None, optional) – Width of the hidden layer of the scalar network, by default n_in_s.

  • scalar_activation (bool, optional) – Whether a SiLU is applied to the output scalars (the hidden blocks of a stack); off for the last block. Default True.

forward(s, v)[source]#

Transform scalar and vector features.

Parameters:

  • s (Tensor) – Scalar features (N, n_in_s).

  • v (Tensor) – Vector features (N, 3, n_in_v).

Returns:

New scalars (N, n_out_s) and vectors (N, 3, n_out_v).

Return type:

tuple of Tensor