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 CUDA or Mosaic TPU 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.

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.

  • 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.

Note

On the CUDA convolution kernel 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.