e3j.utils.options.Layout

class e3j.utils.options.Layout

Memory layout for equivariant arrays.

An equivariant feature space such as "128x0e + 64x1o" is built from irreducible blocks with multiplicity m and irrep dimension 2l + 1 for each degree l.

The layout controls how different blocks are arranged in memory, eventually factorizing the GCD of multiplicities in a separate channels axis.

Values:

• LEADING_CHANNELS – (batch, channels, dim) – channels on an explicit leading axis.

• TRAILING_CHANNELS – (batch, dim, channels) – channels on a trailing axis. Faster on GPU thanks to coalesced memory access.

• E3NN – (batch, channels * dim) – channels folded into the feature dimension, matching e3nn_jax.IrrepsArray.

Note

The CUDA kernels only support channel counts that are powers of two for now.

Methods

parse(key)
parse_value(key)
__getitem__(name)	Return the member matching name.
__iter__()	Return members in definition order.
__len__()	Return the number of members (no aliases)

Attributes

LEADING_CHANNELS
TRAILING_CHANNELS
E3NN