e3j.core.TensorProduct

class e3j.core.TensorProduct(source: tuple[O3Space, O3Space] | tuple[str, str], target: O3Space | str | None, coef: BCOO | None = None, sort: bool = True, config: Config | None = None, layout: str | Layout = 'LEADING_CHANNELS', mode: str | MixingMode = 'OUTER', normalization: str | TensorProductNormalization = 'NONE')

Bilinear tensor products.

This class implements (equivariant) bilinear maps of the form

\[z_i = \sum_{jk} c_{ijk}\, x_j\, y_k,\]

defined by a 3D coefficient array c in sparse BCOO format.

The sparse evaluation algorithm consists of a loop over non-zero coefficients,

  1. a pull-back of x and y by coefficient indices j and k,

  2. a product of x[:,j], y[:,k] and coefficient values c[i,j,k],

  3. accumulation on output coordinates z[:,i].

Dedicated kernels for GPU and TPU should be selected from the environment in the global Config. Otherwise, run:

e3j.config(tensor_product="FUSED_CUDA")         # CUDA
e3j.config(tensor_product="FUSED_MOSAIC_GPU")   # Pallas MGPU (sm_90+)
e3j.config(tensor_product="FUSED_MOSAIC_TPU")   # Pallas MTPU
e3j.config(tensor_product="UNFUSED")            # plain JAX, all platforms

Stack Clebsch Gordan coefficients or use explicitly given ones.

Parameters:

  • source (representations of both inputs.)

  • target (output representation, inferred by default.) – If coefficients are None, then out is interpreted as a momentum and parity filter which should only contain irreducible blocks of multiplicity 1.

  • coef (sparse array, constructed by stacking Clebsch-Gordan) – coefficients by default.

  • sort (whether to sort the output coordinates by grouping momenta and) – parities, True by default.

  • layout (specifies the channel axis, TRAILING_CHANNELS is faster.)

  • mode (specifies how to mix the channels, can take three values:) –

    • OUTER : tensor product of channels in np.outer fashion. Also used to broadcast one operand on the channel axis.

    • INNER : scalar product of channels, summed over after the Clebsch-Gordan tensor product.

    • MAP : channel-wise tensor products. Only useful with trailing channels layout, since leading axes are mapped over by default.

Note

Operations inherit the dtype of their operands. On FUSED_CUDA, the value dtype may be float16, float32 or float64; float16 needs layout="TRAILING_CHANNELS", since the LEADING_CHANNELS kernel reduces with atomicAdd, which has no __half overload.

Methods

__call__(x, y[, coef])

Evaluate bilinear map on pair of inputs.

clebsch_gordan(target, source_1, source_2[, ...])

Stack Clebsch-Gordan coefficients.

infer_target(source[, target, sort])

Infer target representation from Clebsch-Gordan rules.

sort()

Sort irreducible output blocks by degree and parity.

Attributes

coef

Clebsch-Gordan coefficients.

indices

nnz

Number of non-zero coefficients.

nnz_ratio

Density of the coefficient tensor, i.e. nnz over its size.

shape

values

__call__(x: Array, y: Array, coef: Array | None = None) → Array

Evaluate bilinear map on pair of inputs.

static clebsch_gordan(target: O3Space, source_1: O3Space, source_2: O3Space, normalization: TensorProductNormalization = TensorProductNormalization.NONE) → BCOO

Stack Clebsch-Gordan coefficients.

Returns a 3D-array whose dimensions appear in the same order as the arguments.

coef

Clebsch-Gordan coefficients.

Either computed or retrieved from cache. The coef descriptor is writable and may be assigned a value during __init__.

classmethod infer_target(source: tuple[O3Space, O3Space], target: O3Space | None = None, sort: bool = False) → O3Space

Infer target representation from Clebsch-Gordan rules.

The optional target argument should only contain irreducible representations with multiplicity 1, and ValueError will be raised otherwise.

property nnz: int

Number of non-zero coefficients.

property nnz_ratio: float

Density of the coefficient tensor, i.e. nnz over its size.

sort() → TensorProduct

Sort irreducible output blocks by degree and parity.

The ordering agrees with e3nn. It is performed once for all on the coefficient tensor by permuting target coordinates.