Skip to content

Tensor products

Tensor products combine spherical features into new equivariant features. SO3TensorProduct performs unweighted channelwise Gaunt or Clebsch–Gordan contractions. WeightedTensorProduct accepts external weights, input and output representation specifications, and the allowed coupling paths.

Representations can be written as strings or ordered (multiplicity, (degree, parity)) sequences. For example, "2x0e + 3x1o" and [(2, (0, 1)), (3, (1, -1))] both describe two even scalars followed by three odd vectors. Parity is +1 for even (e) and -1 for odd (o); entry order determines the feature blocks and their indices in the coupling instructions.

The following product couples each scalar or vector block of the first input to a scalar second input, preserving the output representation:

from elfes.modules.nn import WeightedTensorProduct

product = WeightedTensorProduct(
    irreps_in1="2x0e + 3x1o",
    irreps_in2="1x0e",
    irreps_out="2x0e + 3x1o",
    instructions=[(0, 0, 0, "uvu", True), (1, 0, 1, "uvu", True)],
)

SO3TensorProduct 

SO3TensorProduct(
    l_max_left: int,
    l_max_right: int,
    l_max_out: int | None = None,
    *,
    kind: Literal["gaunt", "cg"],
    device: device | str | None = None,
    dtype: dtype | None = None,
)

Bases: Module

Unweighted channelwise Gaunt or CG contraction on CUDA.

Inputs use integral-normalized Wikipedia real spherical harmonics, ordered by degree, then m=-l,...,l, with shape [..., (l_max + 1)**2, channels]. Batch shapes and channel counts must agree. Each channel is contracted independently, without weights.

kind="gaunt" computes a pointwise spherical-function product, projected to the output bandlimit. kind="cg" sums all triangle-allowed paths, including odd degree sums, with sqrt(2*lout+1) times OEQ's unit-norm real Wigner-3j coefficients and no additional path normalization.

Supports CUDA float32/float64 and differentiable adjoint contractions through the vendored OEQ JIT runtime.

Parameters:

  • l_max_left (int) –

    Highest degree of the first input.

  • l_max_right (int) –

    Highest degree of the second input.

  • l_max_out (int | None, default: None ) –

    Highest output degree; defaults to the sum of input degrees.

  • kind (Literal['gaunt', 'cg']) –

    Required coefficient choice: "gaunt" or "cg".

  • device (device | str | None, default: None ) –

    Initial module device; move to the input device before use.

  • dtype (dtype | None, default: None ) –

    Floating precision, defaulting to Torch's default dtype.

forward 

forward(left: Tensor, right: Tensor) -> Tensor

Contract the two inputs independently in each channel.

WeightedTensorProduct 

WeightedTensorProduct(
    irreps_in1: str | Sequence[tuple[int, tuple[int, int]]],
    irreps_in2: str | Sequence[tuple[int, tuple[int, int]]],
    irreps_out: str | Sequence[tuple[int, tuple[int, int]]],
    instructions: Sequence[TensorProductInstruction],
)

Bases: Module

Weighted tensor product with an OpenEquivariance-only CUDA implementation.

The provider supports CUDA float32/float64 tensor products whose instructions are all weighted and use one connection mode, uvu or uvw. Representation specifications are converted to private OpenEquivariance metadata; external e3nn is not required. Weights are external and unshared, in canonical instruction/path order (compatible with e3nn weight packing). Component irrep normalization and element path normalization are fixed. Features use Wikipedia real SH; weight packing does not impose e3nn's spherical-harmonic coordinate convention. CPU construction and state loading are supported, but CPU execution is not. Moving the module to CUDA prepares its static schedule so forward is a traceable weight permutation followed by a registered Torch custom op.

Parameters:

  • irreps_in1 (str | Sequence[tuple[int, tuple[int, int]]]) –

    Representation of the first input, as a string such as "2x0e + 3x1o" or an ordered sequence of (multiplicity, (degree, parity)) entries such as [(2, (0, 1)), (3, (1, -1))]. Parity is +1 (even, e) or -1 (odd, o). Entry order defines the feature blocks; within each block, multiplicity precedes the spherical-component axis.

  • irreps_in2 (str | Sequence[tuple[int, tuple[int, int]]]) –

    Representation of the second input, in the same format.

  • irreps_out (str | Sequence[tuple[int, tuple[int, int]]]) –

    Representation of the output, in the same format.

  • instructions (Sequence[TensorProductInstruction]) –

    Ordered (in1_index, in2_index, out_index, mode, has_weight) coupling paths. Indices refer to entries in the corresponding representation specification. All paths must have has_weight=True and share one mode, uvu or uvw.

weight_numel property 

weight_numel: int

Number of external weights expected for each input row.

TensorProductInstruction 

TensorProductInstruction = tuple[int, int, int, str, bool]