e3j.core.Convolution

class e3j.core.Convolution(source: tuple[O3Space, O3Space], target: O3Space | None = None, *, graph_ordering: str | GraphOrdering, layout: str | Layout = Layout.TRAILING_CHANNELS, avg_num_neighbors: float | None = None, normalization: str | TensorProductNormalization = 'SQRT_DIM_OUT', config: Config | None = None)

Equivariant message-passing convolution.

Computes the aggregated message on receiver nodes given by:

\[m_b = \frac 1 N \sum_a (x_a \otimes y_{ab}) \odot s_{ab}\]

where the sum runs over neighbors \(a\) of the receiver node \(b\), \(\otimes\) denotes a tensor product operation, \(\odot\) denotes a scalar mixing, and \(N\) the average number of neighbors.

The plain JAX implementation consists of the following operations:

  1. Gather node features by senders,

  2. Compute the tensor product of sender features with edge features (typically harmonic embeddings),

  3. Mix tensor product outputs with edge scalars (typically MLP of RBF encodings),

  4. Scatter-add messages on receiver nodes.

Optionally, the sum of messages is rescaled by avg_num_neighbors.

Note

When using the CUDA, Mosaic TPU or Mosaic GPU kernels with SENDER graph ordering, the following additional assumptions should hold:

  1. Symmetric graph: each edge (a, b) has its reverse (b, a),

  2. Graded-symmetric edge features: \(y_{ba} = p . y_{ab}\) per slice, where the parity is \(p = (-1)^l\) for harmonic polynomials.

  3. Symmetric scalars: \(s_{ba} = s_{ab}\) (true for distance-based radial functions).

Note

RECEIVER ordering is not yet implemented for the Mosaic TPU kernel.

Note

The Mosaic GPU kernel takes either ordering. Prefer RECEIVER where the edge list can be sorted either way: it aggregates each receiver’s incoming edges directly, so none of the assumptions above apply, and it keeps the backward’s per-edge gather on the narrow node features rather than on the wider output cotangent.

Initialize a Convolution block from parameters.

Parameters:

  • source – Representations of the two tensor-product inputs (node and edge features). The third source space (mixing scalars) is inferred from the former two in the .source attribute.

  • target – Output representation, inferred by default. Passing a target argument enforces a filter on the output irreducible blocks.

  • graph_ordering – Edge ordering for the graph, can be RECEIVER, SENDER or NONE. When edges are sorted by senders, symmetry assumptions on the graph and edge features should hold. Only the unfused path supports NONE for now; the Mosaic TPU path supports only SENDER, every other fused path takes either.

  • layout – Specifies the channel axis, TRAILING_CHANNELS is faster.

  • avg_num_neighbors – If given, messages are divided by this factor.

  • normalization – Normalization of the tensor product’s Clebsch-Gordan coefficients, see e3j.utils.options.TensorProductNormalization.

  • config – Global e3j.utils.config.Config (optional) pointing to the implementation path. The best available option should be automatically selected based on the environment.

Note

Operations inherit the dtype of their operands. On FUSED_CUDA, the value dtype may be float16, float32 or float64: the fused convolution is trailing channels only, so it never hits the atomicAdd float16 restriction.

Methods

__call__(node_features, edge_features, ...)

Return sum of messages on receiver nodes.

Attributes

coef

Return packed Coef4D coefficients for the forward pass.
__call__(node_features: Array, edge_features: Array, edge_scalars: Array, senders: Array, receivers: Array, node_mask: Array | None = None) → Array

Return sum of messages on receiver nodes.

Leading axes should match the number of nodes n0 or number of edges n1 respectively. Feature dimensions num_x, num_y, num_scalars must match the dimensions of the source attribute.

Node features and edge scalars may carry a trailing channel axis of size num_channels, however edge features are expected to not carry one and are broadcast with node features along the channel axis.

Parameters:

  • node_features – array of shape (num_nodes, num_x, num_channels)

  • edge_features – array of shape (num_edges, num_y)

  • edge_scalars – array of shape (num_edges, num_scalars, num_channels)

  • senders – index vector of length num_edges, in bounds [0, num_nodes)

  • receivers – index vector of length num_edges, in bounds [0, num_nodes).

  • node_mask – optional boolean vector of length num_nodes, True for real nodes and False for padding nodes, which must lie at the tail of the graph. Padding edges are also assumed to only connect padding nodes. The Mosaic TPU kernels additionally stop their edge pipeline after the last real edge, so the saving is granular to one edge block.

Note

On the fused convolution kernels the edges must be sorted by the endpoint selected via graph_ordering. The SENDER ordering additionally requires the symmetry assumptions documented on the class.

coef

Return packed Coef4D coefficients for the forward pass.