e3j.core.LinearIndexwise

class e3j.core.LinearIndexwise(source_irreps: str | O3Space, target_irreps: str | O3Space, layout: str | Layout, num_channels: int | None = None)

Index-wise linear mixing, with weights passed explicitly.

Target blocks absent from source produce zero outputs, and source blocks absent from target are discarded.

Methods

__call__(feat, index, weights)

Mix feat with weights gathered by index.

Attributes

blocks

Per-(l, p) blocks of the linear, in weight order.
__call__(feat: Array, index: Array, weights: list[Array]) → Array

Mix feat with weights gathered by index.

Parameters:

  • feat (jax.Array) – (..., C, dim) or (..., dim, C) features depending on layout, (..., dim) without channels.

  • index (jax.Array) – (...)-array of specie indices.

  • weights (list[jax.Array]) – one (I, [C,] m_out, m_in) array per block of blocks.

property blocks: list[LinearBlock]

Per-(l, p) blocks of the linear, in weight order.