Tensor products

Tensor products are the central learnable interaction in e3nn. They combine two representations and project their outer product onto allowed output irreps using Clebsch–Gordan coefficients.

Standard products

FullTensorProduct emits every allowed output path and has no learned weights:

import mlx.core as mx
from e3nn_mlx import o3

tp = o3.FullTensorProduct("2x1o", "3x1o")
left = mx.random.normal((64, tp.irreps_in1.dim))
right = mx.random.normal((64, tp.irreps_in2.dim))
output = tp(left, right)

FullyConnectedTensorProduct connects all compatible multiplicity channels with learned weights:

tp = o3.FullyConnectedTensorProduct(
    "8x0e + 8x1o",
    "1x0e + 1x1o",
    "16x0e + 8x1o",
)
output = tp(left_features, edge_attributes)

Use ElementwiseTensorProduct for aligned channels and TensorSquare when both inputs are the same value.

Symmetry-reduced products

ReducedTensorProducts constructs an orthonormal change-of-basis tensor subject to index permutation symmetries such as "ij=ji" or "ijk=jik=ikj". filter_ir_mid restricts every sequential Clebsch–Gordan coupling, including the final coupling. filter_ir_out restricts the retained final irreps. Intermediate filtering prunes contraction paths and is not equivalent to slicing an already-constructed output.

lmax = 4
bispectrum = o3.ReducedTensorProducts(
    "ijk=jik=ikj",
    i=o3.Irreps.spherical_harmonics(lmax),
    filter_ir_mid=list(o3.Irrep.iterator(lmax=lmax)),
    filter_ir_out=list(o3.Irrep.iterator(lmax=0)),
)

This construction retains scalar symmetric triple contractions suitable for a bispectrum.

General instructions

The general TensorProduct accepts an explicit output representation and instruction list. Each instruction is (input1_index, input2_index, output_index, connection_mode, has_weight); an optional sixth value scales the path before normalization. This interface is powerful but low level—prefer a standard wrapper when it represents the desired connectivity.

Shared and per-sample weights

With shared_weights=True, weights have shape (weight_numel,). With shared_weights=False, external weights have shape (..., weight_numel), which is useful when a radial network predicts a tensor product for every graph edge.

Shared weights remain one-dimensional during execution. MLX broadcasts them inside each contraction rather than materializing one copy per input item. This keeps dense uvw parameter-gradient reductions compact and applies to every weighted TensorProduct connection mode. Batched and singleton-leading external weights still broadcast to the input leading dimensions as required.

Compilation and generated kernels

Compatible float32, rank-two inputs dispatch to specialized Metal kernels on Apple silicon. Other shapes, dtypes, large dense contractions, and unsupported instruction mixtures automatically use general MLX operations. Both paths implement the same contraction and normalization conventions.

Scalar-path kernel selection considers both the static contraction size and the batch size. Small sparse products retain the fused kernel, while dense FullyConnectedTensorProduct workloads use MLX matrix contractions after their measured kernel crossover.