xnn.hybrid.models.bamboo.GETLayer#

class xnn.hybrid.models.bamboo.GETLayer(dim, num_heads, attn_act, is_first=False, is_last=False)[source]#

Bases: Module

One Graph Equivariant Transformer layer (BAMBOO Supplementary A.1).

A GET layer updates the scalar node feature x_i and, except in the last layer, the vector node feature V_i. It first runs the shared multi-head EdgeMultiheadAttention to get a neighbour value v_j and attention weight a_ij per edge, then:

  • forms the scalar message m_i = sum_j a_ij (v_j * d_ij) where d_ij is the per-head radial edge feature, and the vector message u_i = sum_j v_j * e_ij where e_ij = d_ij * r_hat_ij is the equivariant edge vector (both aggregated onto the centre atom);

  • mixes the current vector feature through learned projections and an inner product w_i = <U1 V_i, U2 V_i> (a rotation-invariant scalar), and combines everything into the scalar/vector updates.

The first layer has no incoming V_i (its vector output is just u_i) and the last layer produces no V_i (only the scalar update, which feeds the read-outs). Every layer’s scalar/vector updates are added residually by the parent BAMBOO.

Parameters:

  • dim (int) – Node scalar feature width.

  • num_heads (int) – Number of attention heads.

  • attn_act (str or torch.nn.Module) – Activation applied to the attention logits (BAMBOO uses GELU).

  • is_first (bool, optional) – First layer (no vector input), by default False.

  • is_last (bool, optional) – Last layer (no vector output), by default False.

forward(node_feat, edge_feat, edge_vec, node_vec, center, neighbor, envelope, n_atoms)[source]#

Apply one GET layer.

Parameters:

  • node_feat (Tensor) – Scalar node features (N, dim).

  • edge_feat (Tensor) – Per-head radial edge feature (E, num_heads, dim_per_head).

  • edge_vec (Tensor or None) – Equivariant edge vector (E, 3, num_heads, dim_per_head); unused by the last layer (pass None).

  • node_vec (Tensor or None) – Vector node features (N, 3, dim); None for the first layer.

  • center (Tensor) – Centre (receiver) atom index per edge, shape (E,).

  • neighbor (Tensor) – Neighbour (sender) atom index per edge, shape (E,).

  • envelope (Tensor) – Radial cutoff weight per edge, shape (E,).

  • n_atoms (int) – Number of atoms N (scatter dimension size).

Returns:

The scalar update delta_x of shape (N, dim) and, unless this is the last layer, the vector update/output of shape (N, 3, dim) (None for the last layer).

Return type:

tuple of (Tensor, Tensor or None)

On this page
Show Source