Linear operations

These layers mix feature channels while preserving the relevant symmetry. SO3Linear acts within each spherical degree; SO2Linear mixes channels in an m-major layout. ComplexLinear and ComplexLinear3M implement the same complex linear map using different real-valued matrix operations.

The frame and layout transforms describe how spherical features are prepared for local SO(2) operations.

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].

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].

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].