elfes.modules.nn

Reusable operations for constructing and transforming network features.

Irreps

Bases: tuple

slices

slices()

List of slices corresponding to indices for each irrep.

Examples:

Irreps('2x0e + 1e').slices() [slice(0, 2, None), slice(2, 5, None)]

count

count(ir: Irrep) -> int

:returns: total multiplicity of ir.

mul

__mul__(other) -> Irreps

(Irreps('2x1e') * 3).simplify() 6x1e

rmul

__rmul__(other) -> Irreps

2 * Irreps('0e + 1e') 1x0e+1x1e+1x0e+1x1e

ComplexLinear

ComplexLinear(
    in_channels: int,
    out_channels: int,
    *,
    device: device | str | None = None,
    dtype: dtype | None = None,
)

Bases: Module

Apply a bias-free complex matrix to real/imaginary channel pairs.

For \(z=x_0+i x_1\) and \(W=A+iB\), compute \(y=Wz\) as

\[ y_0=Ax_0-Bx_1,\qquad y_1=Bx_0+Ax_1. \]

Inputs are real tensors shaped [..., 2, C_in]; axis -2 holds the real and imaginary parts. Outputs use the same layout, [..., 2, C_out]. The weight stores \((A,B)\) with shape [2, C_out, C_in].

Each call builds the real block matrix \([[A,-B],[B,A]]\) and applies one real linear map to the flattened pair. Only the final real/imaginary output is materialized, without storing the four individual products.

The map commutes with any common phase multiplication \(z\mapsto e^{i\phi}z\). This makes it suitable for copies of a nonzero SO(2)/U(1) mode, without depending on the mode index or the origin of the channels. A nonzero additive bias would break this phase equivariance and is not supported. Real invariant channels, such as the \(m=0\) block of SO2Linear, use an ordinary real nn.Linear instead.

Parameters:

in_channels
(int) –

Number of complex input channels.
out_channels
(int) –

Number of complex output channels.
device
(device | str | None, default: None ) –

Device on which to create the weight.
dtype
(dtype | None, default: None ) –

Real dtype used to store weights and channel pairs.

forward

forward(complex_tensor: Tensor) -> Tensor

Parameters:

complex_tensor
(Tensor) –

[..., 2, C_in].

Returns:

output_complex_tensor ( Tensor ) –

[..., 2, C_out].

ComplexLinear3M

ComplexLinear3M(
    in_channels: int,
    out_channels: int,
    *,
    device: device | str | None = None,
    dtype: dtype | None = None,
)

Bases: Module

Apply a bias-free complex linear map using three real matrix products.

This module has the same API and parameter layout as ComplexLinear. For the complex input \(z=x_0+i x_1\) and complex weight \(W=A+iB\), it computes

\[ p=Ax_0,\qquad q=Bx_1,\qquad r=(A+B)(x_0+x_1), \]

and recovers the complex product as

\[ \operatorname{Re}(Wz)=p-q,\qquad \operatorname{Im}(Wz)=r-p-q. \]

Inputs have shape [..., 2, C_in] and outputs have shape [..., 2, C_out], with axis -2 containing \((x_0,x_1)\). The weight has shape [2, C_out, C_in] and stores \((A,B)\).

Parameters:

in_channels
(int) –

Number of complex input channels.
out_channels
(int) –

Number of complex output channels.
device
(device | str | None, default: None ) –

Device on which to create the weight.
dtype
(dtype | None, default: None ) –

Dtype with which to create the real weight.

forward

forward(complex_tensor: Tensor) -> Tensor

Parameters:

complex_tensor
(Tensor) –

[..., 2, C_in].

Returns:

output_complex_tensor ( Tensor ) –

[..., 2, C_out].

IrrepLayoutTransform

IrrepLayoutTransform(
    l_max: int, m_max: int | None = None, *, device: device | str | None = None
)

Bases: Module

Reorder SO(3) spherical tensors between l-major and m-major layouts.

Both layouts contain the same components with 0 <= l <= l_max and |m| <= min(l,m_max). The m_max cutoff does not remove any l block; it removes the components with |m| > m_max from every higher-l block. Thus with l_max=3 and m_max=2, the l=3 block retains m=-2,...,+2 but omits m=-3 and m=+3.

For this example, the l-major layout, written as (l,m), is:

l=0: [(0,0)] l=1: [(1,-1), (1,0), (1,+1)] l=2: [(2,-2), (2,-1), (2,0), (2,+1), (2,+2)] l=3: [(3,-2), (3,-1), (3,0), (3,+1), (3,+2)]

The corresponding m-major layout, written as (m,l), is:

m=0: [(0,0), (0,1), (0,2), (0,3)] m=+1: [(+1,1), (+1,2), (+1,3)] m=-1: [(-1,1), (-1,2), (-1,3)] m=+2: [(+2,2), (+2,3)] m=-2: [(-2,2), (-2,3)]

This transform does not truncate m components. It accepts tensors already truncated to |m| <= m_max and only reorders them. to_m_major and to_l_major are exact inverses.

Parameters:

l_max
(int) –

Highest non-negative l.
m_max
(int | None, default: None ) –

Highest retained |m|. Defaults to l_max.
device
(device | str | None, default: None ) –

Initial device of the permutation indices.

Shape

Input: [..., K, C]
Output: [..., K, C]

Here K = (m_max+1)**2+(l_max-m_max)*(2*m_max+1).

Attributes:

m_sizes –

Number of retained l values for each non-negative m.
n_components –

Number of retained spherical components.

to_m_major

to_m_major(l_major_tensor: Tensor) -> Tensor

Parameters:

l_major_tensor
(Tensor) –

[..., K, C].

Returns:

so2_tensor ( Tensor ) –

[..., K, C].

to_l_major

to_l_major(so2_tensor: Tensor) -> Tensor

Parameters:

so2_tensor
(Tensor) –

[..., K, C].

Returns:

l_major_tensor ( Tensor ) –

[..., K, C].

EquivariantRMSNorm

EquivariantRMSNorm(
    l_max: int,
    channels: int,
    *,
    eps: float = 1e-05,
    learnable: bool = True,
    center: bool = True,
    weighting: Literal["l_balanced", "component"] = "l_balanced",
    device: device | str | None = None,
    dtype: dtype | None = None,
)

Bases: Module

Normalize complete SO(3) tensors with one RMS per sample.

Inputs contain one complete real SO(3) irrep for every l = 0, ..., l_max and use the l-major shape (..., (l_max + 1) ** 2, channels). When center is true, the scalar channels are first centered as

\[ \mu=\frac{1}{C}\sum_c x_{00c}, \qquad \widetilde{x}_{lmc}=x_{lmc}-\delta_{l0}\mu. \]

The default l_balanced weighting gives every l block equal weight:

\[ q=\frac{1}{C}\sum_c\frac{1}{l_{\max}+1} \sum_{l=0}^{l_{\max}}\frac{1}{2l+1} \sum_{m=-l}^{l}\widetilde{x}_{lmc}^2. \]

component weighting instead averages all (l, m, c) entries uniformly. The normalized tensor shares the sample-dependent scale \((q+\varepsilon)^{-1/2}\). When learnable is true, the result is multiplied by a learned weight[l, c] shared across m; centered scalar channels also receive a learned bias[c]:

\[ y_{lmc}=\gamma_{lc}\frac{\widetilde{x}_{lmc}} {\sqrt{q+\varepsilon}}+\delta_{l0}\beta_c. \]

Sharing parameters across m and restricting the bias to l = 0 makes the operation commute with rotations. Contiguous CUDA float16, bfloat16, and float32 inputs use fused Triton forward and first-order backward kernels; higher-order gradients and other inputs use differentiable PyTorch expressions.

Parameters:

l_max
(int) –

Highest included non-negative l.
channels
(int) –

Number of channels shared by every spherical component.
eps
(float, default: 1e-05 ) –

Positive value added before the reciprocal square root.
learnable
(bool, default: True ) –

Learn the equivariant scale and, when centering, scalar bias.
center
(bool, default: True ) –

Subtract the mean of the l = 0 channels for every sample.
weighting
(Literal['l_balanced', 'component'], default: 'l_balanced' ) –

Give each l block or each spherical component equal weight.
device
(device | str | None, default: None ) –

Initial device of parameters and buffers.
dtype
(dtype | None, default: None ) –

Initial floating dtype of parameters and weights.

Attributes:

weight (Tensor | None) –

Learned scale with shape (l_max + 1, channels), or None.
bias (Tensor | None) –

Learned scalar bias with shape (channels,), or None.

forward

forward(so3_tensor: Tensor) -> Tensor

Parameters:

so3_tensor
(Tensor) –

[..., (l_max+1)**2, C].

Returns:

normalized_so3_tensor ( Tensor ) –

[..., (l_max+1)**2, C].

GaussianBasis

GaussianBasis(cutoff: float, num_basis: int, *, width_scale: float = 2.0)

Bases: Module

Fixed Gaussian basis with uniformly spaced centers.

For \(n = 0, \ldots, N - 1\), define

\[ \begin{aligned} \mu_n &= \frac{n r_\mathrm{c}}{N - 1}, \\ \Delta &= \frac{r_\mathrm{c}}{N - 1}, \\ \sigma &= s\Delta, \end{aligned} \]

The basis functions are

\[ G_n(r) = \exp\left[-\frac{1}{2}\left(\frac{r - \mu_n}{\sigma}\right)^2\right]. \]

The centers and widths are fixed rather than learned. Gaussian functions do not vanish at \(r_\mathrm{c}\), so a separate cutoff remains necessary when edge messages must go smoothly to zero.

Parameters:

cutoff
(float) –

Cutoff radius \(r_\mathrm{c}\) and location of the final center.
num_basis
(int) –

Number of Gaussian functions \(N\); must be at least two.
width_scale
(float, default: 2.0 ) –

Width scale \(s\) relative to the center spacing.

LearnableSphericalBesselBasis

LearnableSphericalBesselBasis(cutoff: float, num_basis: int)

Bases: _SphericalBesselBasis

Learnable zero-order spherical Bessel radial basis.

This basis has the same initial functions as :class:SphericalBesselBasis, but optimizes their dimensionless frequencies, initialized to \(n\pi\) for \(n = 1, \ldots, N\).

Parameters:

cutoff
(float) –

Cutoff radius \(r_\mathrm{c}\).
num_basis
(int) –

Number of basis functions \(N\).

PolynomialEnvelope

PolynomialEnvelope(cutoff: float, power: int = 6)

Bases: Module

DimeNet polynomial envelope with two vanishing derivatives.

With \(x = r / r_\mathrm{c}\) and positive integer power \(p\), the cutoff is

\[ E(r) = \begin{cases} 1 - \dfrac{(p+1)(p+2)}{2}x^p + p(p+2)x^{p+1} - \dfrac{p(p+1)}{2}x^{p+2}, & r < r_\mathrm{c}, \\ 0, & r \ge r_\mathrm{c}. \end{cases} \]

Its value, slope, and curvature all reach zero at \(r_\mathrm{c}\), giving a smooth transition to zero. The default \(p = 6\) is the common NequIP-style choice.

Parameters:

cutoff
(float) –

Cutoff radius \(r_\mathrm{c}\).
power
(int, default: 6 ) –

Polynomial power \(p\).

SphericalBesselBasis

SphericalBesselBasis(cutoff: float, num_basis: int)

Bases: _SphericalBesselBasis

Fixed orthonormal zero-order spherical Bessel radial basis.

For cutoff radius \(r_\mathrm{c}\) and \(n = 1, \ldots, N\), the basis functions are

\[ B_n(r) = \sqrt{\frac{2}{r_\mathrm{c}}} \frac{\sin(n\pi r / r_\mathrm{c})}{r}. \]

They are orthonormal on \([0, r_\mathrm{c}]\) with radial measure \(r^2\,\mathrm{d}r\). At the origin, the implementation uses the analytic limit

\[ B_n(0) = \sqrt{\frac{2}{r_\mathrm{c}}}\frac{n\pi}{r_\mathrm{c}}. \]

The frequencies \(n\pi\) are fixed.

Parameters:

cutoff
(float) –

Cutoff radius \(r_\mathrm{c}\).
num_basis
(int) –

Number of basis functions \(N\).

S2GridTransform

S2GridTransform(
    l_max: int,
    m_max: int | None = None,
    grid_shape: tuple[int, int] | None = None,
    *,
    component_layout: Literal["l_major", "m_major"] = "l_major",
    device: device | str | None = None,
    dtype: dtype | None = None,
)

Bases: Module

Transform band-limited real spherical functions to and from an S2 grid.

The component basis is the normalized Wikipedia real spherical-harmonic basis used throughout ELFES. Setting m_max < l_max keeps only |m| <= m_max in every \(l\) block. Components use either l_major layout, with consecutive \(l\) blocks ordered by m = -l, ..., l, or m_major layout, with the \(m=0\) block followed by the positive and negative blocks for each \(|m|\) and \(l\) increasing within every block.

The equiangular grid uses polar angles beta_i = pi * (i + 1/2) / n_polar and azimuthal angles alpha_j = 2 * pi * j / n_azimuthal. Its Cartesian directions are (sin(beta) * cos(alpha), sin(beta) * sin(alpha), cos(beta)). The inverse uses the Kostelec–Rockmore quadrature for this grid.

Parameters:

l_max
(int) –

Highest non-negative \(l\).
m_max
(int | None, default: None ) –

Highest retained \(|m|\). Defaults to l_max.
grid_shape
(tuple[int, int] | None, default: None ) –

Number of polar and azimuthal samples. Defaults to (2 * (l_max + 1), 2 * m_max + 1).
component_layout
(Literal['l_major', 'm_major'], default: 'l_major' ) –

Layout of the spherical components.
device
(device | str | None, default: None ) –

Initial buffer device.
dtype
(dtype | None, default: None ) –

Initial floating dtype. Defaults to PyTorch's default dtype.

Shape

Components: (..., n_components, channels)
Grid values: (..., n_polar, n_azimuthal, channels)

Here n_components = sum(2 * min(l, m_max) + 1 for l in range(l_max + 1)).

forward

forward(components: Tensor) -> Tensor

Alias for to_grid.

to_grid

to_grid(components: Tensor) -> Tensor

Evaluate spherical components on the configured grid.

from_grid

from_grid(values: Tensor) -> Tensor

Project values on the configured grid back to components.

SO2Linear

SO2Linear(
    in_channels: int,
    out_channels: int,
    l_max: int,
    m_max: int | None = None,
    *,
    bias: bool = True,
    algorithm: Literal["direct", "3m"] = "direct",
    device: device | str | None = None,
    dtype: dtype | None = None,
)

Bases: Module

Apply an SO(2)-equivariant linear map to SO(3) spherical tensors.

The input and output use m-major layout. They first store every \(m=0\) component in increasing \(l\) order. For each \(m>0\), they then store every \((l,+m)\) component followed by every \((l,-m)\) component, both in increasing \(l\) order. The operation views these tensors after restricting rotations to the SO(2) subgroup about the local z axis. It preserves \(|m|\) while freely mixing all \((l,c)\) copies within each fixed \(|m|\).

For \(m=0\), the components are invariant under SO(2), so the block is an ordinary real linear map over the combined \((l,c)\) axis:

\[ y_{l0c_o} = \sum_{l',c_i} W^{(0)}_{l c_o,l' c_i}x_{l'0c_i} +b_{l c_o}. \]

For each \(m>0\), define one complex value from the two real components

\[ z_{lmc}=x_{l,+m,c}+i x_{l,-m,c}. \]

The corresponding block is the complex linear map

\[ z'_{lmc_o} = \sum_{l',c_i} W^{(m)}_{l c_o,l' c_i}z_{l'mc_i}, \]

implemented by either ComplexLinear or ComplexLinear3M on real tensors. A bias is allowed only in the invariant \(m=0\) block.

This module is not SO(3)-equivariant by itself because it may mix different \(l\). It is intended for tensors already expressed in a local frame; the composition of rotation to that frame, this SO(2)-equivariant map, and rotation back to the global frame can be SO(3)-equivariant.

Parameters:

in_channels
(int) –

Number of input channels for every spherical component.
out_channels
(int) –

Number of output channels for every spherical component.
l_max
(int) –

Highest non-negative \(l\).
m_max
(int | None, default: None ) –

Highest retained \(|m|\). Defaults to l_max.
bias
(bool, default: True ) –

Whether the \(m=0\) block has a bias.
algorithm
(Literal['direct', '3m'], default: 'direct' ) –

Real multiplication algorithm for every \(m>0\) block. Defaults to direct, the block-weight ComplexLinear.
device
(device | str | None, default: None ) –

Initial device of parameters and layout indices.
dtype
(dtype | None, default: None ) –

Initial floating dtype of the parameters.

Shape

Input: [..., K, C_in] in m-major layout
Output: [..., K, C_out] in m-major layout

Here K = (m_max+1)**2+(l_max-m_max)*(2*m_max+1).

Attributes:

m_sizes –

Number of retained \(l\) values for each non-negative \(m\).
n_components –

Number of retained spherical components.

forward

forward(so2_tensor: Tensor) -> Tensor

The K axis uses an m-major layout: all m=0 components come first, followed for each m>0 by all +m components and then all -m components, with l increasing inside each group. Written as (m,l), the order for l_max=3 and m_max=2 is:

m=0: [(0,0), (0,1), (0,2), (0,3)] m=+1: [(+1,1), (+1,2), (+1,3)] m=-1: [(-1,1), (-1,2), (-1,3)] m=+2: [(+2,2), (+2,3)] m=-2: [(-2,2), (-2,3)]

Parameters:

so2_tensor
(Tensor) –

[..., K, C_in], where K = (m_max+1)**2+(l_max-m_max)*(2*m_max+1).

Returns:

output_so2_tensor ( Tensor ) –

[..., K, C_out].

SO3Linear

SO3Linear(
    in_channels: int,
    out_channels: int,
    l_max: int,
    *,
    bias: bool = True,
    device: device | str | None = None,
    dtype: dtype | None = None,
)

Bases: Module

Mix channels independently within each spherical l block.

For every l = 0, ..., l_max and m = -l, ..., l, the layer applies

\[ y_{l m c_o} = \sum_{c_i=1}^{C_\mathrm{in}} W^{(l)}_{c_o c_i} x_{l m c_i} + \delta_{l0} b_{c_o}. \]

Each \(W^{(l)}\) is shared by every \(m\) in the same \(l\) block, so the map commutes with SO(3) rotations. The bias \(b\) is present only when bias is true and acts only on \(l = 0\). Inputs use the complete \(l\)-major layout [..., (l_max+1)**2, C_in]; outputs have shape [..., (l_max+1)**2, C_out].

Parameters:

in_channels
(int) –

Number of input channels shared by every \(l\).
out_channels
(int) –

Number of output channels shared by every \(l\).
l_max
(int) –

Highest included non-negative \(l\).
bias
(bool, default: True ) –

Add a learned scalar bias when true.
device
(device | str | None, default: None ) –

Device on which to create parameters.
dtype
(dtype | None, default: None ) –

Dtype with which to create parameters.

forward

forward(so3_tensor: Tensor) -> Tensor

Parameters:

so3_tensor
(Tensor) –

[..., (l_max+1)**2, C_in].

Returns:

output_so3_tensor ( Tensor ) –

[..., (l_max+1)**2, C_out].

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: Irreps,
    irreps_in2: Irreps,
    irreps_out: Irreps,
    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. Irreps metadata comes from vendored e3nn-lite, not external e3nn. 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.

weight_numel `property`

weight_numel: int

Number of external weights expected for each input row.

build_local_frame

build_local_frame(z_direction: Tensor, x_reference: Tensor) -> Tensor

Build proper global-to-local frames from z directions and x references.

Each nonzero z_direction is normalized to form the local z axis. The corresponding x_reference is projected onto its orthogonal plane and normalized to form the local x axis; the local y axis completes a right-handed frame. A stable Cartesian reference is used when the supplied reference is nearly parallel to z.

Parameters:

z_direction
(Tensor) –

Nonzero Cartesian directions with shape (..., 3).
x_reference
(Tensor) –

Cartesian reference directions broadcastable with z_direction and on the same device with the same dtype. The projected direction fixes the remaining SO(2) gauge freedom around z.

Returns:

Tensor –

Proper rotation matrices with shape (..., 3, 3), whose rows are the
Tensor –

local x, y, and z axes in global coordinates. Therefore the matrices
Tensor –

map global Cartesian coordinates into the constructed local frames.

to_global_frame

to_global_frame(
    so2_tensor: Tensor,
    wigner_matrices: Tensor,
    l_max: int,
    m_max: int | None = None,
    *,
    so2_layout: Literal["l_major", "m_major"] = "l_major",
) -> Tensor

Zero-fill local features and rotate them into the global frame.

Restore l-major ordering, fill omitted m components with zeros, and apply the transpose Wigner rotation. This is the adjoint of to_local_frame: without truncation it undoes that transform, but with truncation it cannot recover discarded components. Channels are not mixed.

Parameters:

so2_tensor
(Tensor) –

Local features shaped (..., K, channels) in so2_layout, where K = sum(2*min(l, M)+1 for l in range(l_max+1)) and M = l_max if m_max is None else m_max.
wigner_matrices
(Tensor) –

The same global-to-local Wigner matrices used by to_local_frame; do not transpose them yourself. Shape is (..., sum((2*l+1)**2 for l in range(l_max+1))); leading dimensions, dtype, and device must match so2_tensor.
l_max
(int) –

Highest included non-negative degree.
m_max
(int | None, default: None ) –

Highest absolute order present in the local input; None means all orders are present.
so2_layout
(Literal['l_major', 'm_major'], default: 'l_major' ) –

Local input layout, as described by to_local_frame. This must match the layout used to construct so2_tensor.

Returns:

Tensor –

Complete global so3_tensor shaped (..., (l_max+1)**2, channels)
Tensor –

in l-major order. Leading dimensions, channels, dtype, and device
Tensor –

are preserved.

to_local_frame

to_local_frame(
    so3_tensor: Tensor,
    wigner_matrices: Tensor,
    l_max: int,
    m_max: int | None = None,
    *,
    so2_layout: Literal["l_major", "m_major"] = "l_major",
) -> Tensor

Rotate global SO(3) features into local frames and optionally truncate m.

Apply the Wigner rotation, retain components with \(|m|\leq m_{\max}\), and pack the local features in so2_layout. The complete global input is always l-major. Channels are not mixed by this operation.

Parameters:

so3_tensor
(Tensor) –

Complete global features shaped (..., (l_max+1)**2, channels) in l-major order.
wigner_matrices
(Tensor) –

Flattened matrices from wigner_D(l_max, frames), where frames maps global Cartesian coordinates to local ones, as returned by build_local_frame. Shape is (..., sum((2*l+1)**2 for l in range(l_max+1))); leading dimensions, dtype, and device must match so3_tensor.
l_max
(int) –

Highest included non-negative degree.
m_max
(int | None, default: None ) –

Highest retained absolute order; None retains all orders.
so2_layout
(Literal['l_major', 'm_major'], default: 'l_major' ) –

Local output layout. l_major orders by ascending l, then ascending m within each l block. m_major orders blocks as m=0,+1,-1,+2,-2,..., with ascending l inside each block.

Returns:

Tensor –

Local so2_tensor shaped (..., K, channels) in so2_layout, where
Tensor –

K = sum(2*min(l, M)+1 for l in range(l_max+1)) and
Tensor –

M = l_max if m_max is None else m_max. Leading dimensions,
Tensor –

channels, dtype, and device are preserved.

Use to_global_frame with the same matrices and settings to rotate back. If orders were truncated, their original values cannot be recovered.