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 | TPMode = 'OUTER', normalization: str | TPNormalization = 'NONE')

Bilinear tensor products.

This class implements (equivariant) bilinear maps of the form:

z[:,i] = Σⱼₖ c[i,j,k] * 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
e3j.config(tensor_product="FUSED_MOSAIC_TPU")   # Pallas MTPU
e3j.config(tensor_product="SPARSE")             # plain JAX, all platforms

Methods

__call__(x, y[, coef])

Evaluate bilinear map on pair of inputs.

__init__(source, target[, coef, sort, ...])

Stack Clebsch Gordan coefficients or use explicitly given ones.

aggregate(values[, layout])

Aggregate a values-like vector on output coordinates.

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

Stack Clebsch-Gordan coefficients.

dense_eval(x, y, coef)

Evaluate bilinear map on pair of inputs.

fused_eval(x, y, coef)

Evaluate bilinear map on pair of inputs.

infer_target(source[, target, sort])

Infer target representation from Clebsch-Gordan rules.

mtpu_eval(x, y, coef)

Evaluate bilinear map via the Pallas Mosaic TPU kernel.

sort()

Sort irreducible output blocks.

sparse_eval(x, y, coef)

Evaluate bilinear map on pair of inputs.

Attributes

aggregation_method

Cache self._aggregation_method with default.

coef

A writable cache descriptor.

indices

is_dense

is_fused

is_mtpu

nnz

nnz_ratio

shape

target_matrix

Return aggregation matrix of shape (nnz, dim_out).

values

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

Evaluate bilinear map on pair of inputs.

__init__(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 | TPMode = 'OUTER', normalization: str | TPNormalization = 'NONE')

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.

aggregate(values, layout: Layout | str = 'E3NN')

Aggregate a values-like vector on output coordinates.

Parameters:

  • values (np.ndarray) – a vector of length self.nnz, e.g. the product of CG coefficients with input coordinates.

  • layout – array layout for scatter-based aggregation.

Returns:

a vector of shape target.dim.

Return type:

np.ndarray

property aggregation_method: Aggregation

Cache self._aggregation_method with default.

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

Stack Clebsch-Gordan coefficients.

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

dense_eval(x: Array, y: Array, coef: Array) → Array

Evaluate bilinear map on pair of inputs.

fused_eval(x: Array, y: Array, coef: Array) → Array

Evaluate bilinear map on pair of inputs.

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.

mtpu_eval(x: Array, y: Array, coef: Array) → Array

Evaluate bilinear map via the Pallas Mosaic TPU kernel.

sort() → TensorProduct

Sort irreducible output blocks.

In contrast with e3nn, the reordering of output coordinates is performed once and for all on the full coefficient tensor, avoiding the associated overhead at evaluation time.

sparse_eval(x: Array, y: Array, coef: Array) → Array

Evaluate bilinear map on pair of inputs.

property target_matrix: BCOO | Array

Return aggregation matrix of shape (nnz, dim_out).

By default, the aggregation matrix is a dense array. To get a sparse array, set the environment variable E3J_AGGREGATION_METHOD to “sparse”.

Example

Number of coefficients above each output coordinate:

>>> num_coefs = jnp.ones(op.nnz) @ op.target_matrix