Feat: mosaic tpu tensor product

[olivier-peltre]

Adds forward and backward tensor product Pallas kernels for TPU.

Some context

The tensor product consists of a bilinear contraction of (x, y) with a sparse 3D array of coefficients (see core.TensorProduct):

$$ z_i = \sum_{j,k} c_{ijk} \cdot x_j \cdot y_k $$

where $c$ is passed in sparse COO format (Clebsch-Gordan coefficients) sorted lexicographically by output index first. Sparse e3j kernels therefore loop over coefficients to load $x_j$ and $y_k$ from buffers (here VMEM, 2-stage pipelining).

Baseline

In contrast the e3nn baseline uses a block-sparse for loop of dense einsums (see tensor_products.py)

zs = []
for l1, l2 in zip(degrees_lhs, degrees_rhs): 
    for l0 in range(abs(l1 - l2), l1 + l2): 
        # NOTE: (l0, l1, l2) indexes a nonzero 3D-block of the full
        #       Clebsch-Gordan tensor. Each axis has size 2l+1. 
        #       The `coef` array is dense although it has a fine-grained
        #       sparsity pattern (nnz scales quadratically in a cubic block).   
        coef = e3nn.clebsch_gordan(l0, l1, l2)
        x_l1 = jnp.take(x, slice_l1)
        y_l2 = jnp.take(y, slice_l2)
        z_l0 = jnp.einsum(
            "ijk,..cj,...ck -> ...ci", coef, x, y
        )
        zs.append((l0, z_l0))

Channel mixing

Each irreducible block (indexed by degree $l$, of dimension $2l + 1$) has a multiplicity, typically an MLIP's number of channels (32, 64, 128...). Different implementations may use different layouts, i.e.

• Layout.E3NN: the string "128x0e + 128x1o" denotes 128 scalars followed by 128 vectors (flattened)
• Layout.TRAILING_CHANNELS (our kernels): multiplicities are factored in a trailing channel axis, so "128x0e + 128x1o" is represented as (1+3, 128) shaped arrays.
• Layout.LEADING_CHANNELS: channel axis before the trailing feature axis, usually less efficient.

The way channel contents are mixed only needs to be bilinear, there are otherwise no constraints coming from equivariance. The channel mixing mode thus depends on downstream settings:

• TPMode.MAP : channel-wise bilinear couplings (x and y have equal channel dimensions)
• TPMode.OUTER : broadcast y with x along the channel axis (y has 1 channel) as a particular
  case of np.outer-like mixing
• TPMode.INNER : sum along the channel axis after the bilinear coupling (x and y have equal channel dimensions)

Note: INNER is the dual of OUTER, which means it occurs during the backward pass of an outer coupling (leftout from this PR's forward kernel).

Benchmarks (TPUv4)

The forward and backward kernels reach near-HBM throughput: ~1 TB/s with 128 channels and l_max=3, close to the 1.2 TB/s advertized IIUC.

Forward:

Backward:

A follow up PR by @armandpicard should address the relative performance gap at low multiplicities by packing more copies in a single buffer.