eLinear

class eLinear(in_rep, out_rep, bias=True, init_scheme='xavier_normal', basis_expansion_scheme='isotypic_expansion')

Bases: Linear

Parameterize a \(\mathbb{G}\)-equivariant linear map with optional invariant bias.

The layer learns coefficients over \(\operatorname{Hom}_\mathbb{G}(\rho_{\text{in}}, \rho_{\text{out}})\), synthesizing a dense weight matrix \(\mathbf{W}\) satisfying:

\[\rho_{\text{out}}(g) \mathbf{W} = \mathbf{W} \rho_{\text{in}}(g) \quad \forall g \in \mathbb{G}\]

If bias=True, the bias vector \(\mathbf{b}\) is constrained to the invariant subspace:

\[\rho_{\text{out}}(g) \mathbf{b} = \mathbf{b} \quad \forall g \in \mathbb{G}\]

Note

Runtime behavior depends on mode. In training mode (model.train()), the constrained dense tensors are recomputed every forward pass, which is correct for gradient updates but slower. In inference mode (model.eval()), the expanded dense weight (and optional invariant bias) are cached and reused until parameters change or invalidate_cache() is called, which is faster. With the cache active, forward() is computationally equivalent to a symmetry-agnostic Linear with fixed dense weight and bias.

homo_basis

Handler exposing the equivariant basis and metadata.

Type:

GroupHomomorphismBasis

bias_module

Optional module handling the invariant bias.

Type:

InvariantBias | None

Initialize the equivariant layer.

Parameters:

• in_rep (Representation) – Representation \(\rho_{\text{in}}\) describing how inputs transform.

• out_rep (Representation) – Representation \(\rho_{\text{out}}\) describing how outputs transform.

• bias (bool, optional) – Enables the invariant bias if the trivial irrep is present in out_rep. Default: True.

• init_scheme (str | None, optional) – Initialization method passed to initialize_params(). Use None to skip initialization. Default: "xavier_normal".

• basis_expansion_scheme (str, optional) – Strategy for materializing the basis ("isotypic_expansion" or "memory_heavy"). Default: "isotypic_expansion".

Raises:

ValueError – If \(\dim(\mathrm{Hom}_{\mathbb{G}}(\rho_{\text{in}}, \rho_{\text{out}})) = 0\).

property bias: Tensor | None

Invariant bias from InvariantBias (None if disabled).

expand_weight()

Return the dense equivariant weight, caching it outside training.

Returns:

Dense matrix of shape (out_rep.size, in_rep.size).

Return type:

torch.Tensor

invalidate_cache()

Clear cached expansions and mark them stale.

Return type:

None

load_state_dict(state_dict, strict=True)

Load parameters and invalidate cached expanded tensors.

Parameters:

strict (bool)

reset_parameters(scheme='xavier_normal')

Reset all trainable parameters.

Parameters:

scheme (str) – Initialization scheme ("xavier_normal", "xavier_uniform", "kaiming_normal", or "kaiming_uniform").

train(mode=True)

Switch mode and keep cached expanded tensors consistent.

Parameters:

mode (bool)

property weight: Tensor

Dense equivariant weight; recomputed in train, cached in eval.