GroupHomomorphismBasis#

class GroupHomomorphismBasis(in_rep, out_rep, basis_expansion='isotypic_expansion')[source]#

Bases: Module

Basis handler for \(\operatorname{Hom}_\mathbb{G}(\rho_{\mathcal{X}}, \rho_{\mathcal{Y}})\).

This module provides:

  1. Synthesis of equivariant maps from basis coefficients.

  2. Orthogonal projection of dense matrices onto the equivariant subspace.

Using the isotypic decomposition (isotypic_decomp_rep()) of \(\rho_{\mathcal{X}}\) and \(\rho_{\mathcal{Y}}\), each shared irrep type \(k\) contributes a block with structure

\[\mathbf{A}^{(k)}_{o,i} = \sum_{s=1}^{S_k}\theta^{(k)}_{s,o,i}\,\mathbf{\Psi}^{(k)}_s,\]

where \(\{\mathbf{\Psi}^{(k)}_s\}_{s=1}^{S_k}\) is a basis of \(\operatorname{End}_\mathbb{G}(\hat{\rho}_k)\). This induces \(m^{\text{out}}_k m^{\text{in}}_k S_k\) degrees of freedom for type \(k\).

Core utilities:

  • forward(): map basis coefficients to a dense equivariant matrix \(\mathbf{W}\in\operatorname{Hom}_\mathbb{G}(\rho_{\mathcal{X}},\rho_{\mathcal{Y}})\).

  • orthogonal_projection(): project any dense matrix onto \(\operatorname{Hom}_\mathbb{G}(\rho_{\mathcal{X}},\rho_{\mathcal{Y}})\) in Frobenius norm.

  • initialize_params(): sample valid initialization either in coefficient space or as a dense equivariant matrix.

Examples

Build and synthesize an equivariant matrix with forward().

>>> from escnn.group import CyclicGroup
>>> G = CyclicGroup(4)
>>> rep = direct_sum([G.regular_representation] * 2)
>>> basis = GroupHomomorphismBasis(rep, rep)
>>> w_dof = torch.randn(basis.dim)
>>> W = basis(w_dof)
>>> # W satisfies rho(g) @ W == W @ rho(g) for all g in G

Project an unconstrained matrix with orthogonal_projection().

>>> W0 = torch.randn(rep.size, rep.size)
>>> W_proj = basis.orthogonal_projection(W0)
>>> # W_proj is the closest equivariant matrix to W0 in Frobenius norm

Sample initialization with initialize_params().

>>> w0 = basis.initialize_params("xavier_normal", return_dense=False)
>>> W0 = basis.initialize_params("xavier_normal", return_dense=True)
>>> # both parameterize equivariant maps
G#

Symmetry group shared by in_rep and out_rep.

Type:

Group

in_rep#

Input representation \(\rho_{\mathcal{X}}\) rewritten in an isotypic basis.

Type:

Representation

out_rep#

Output representation \(\rho_{\mathcal{Y}}\) rewritten in an isotypic basis.

Type:

Representation

basis_expansion#

Strategy (“memory_heavy” or “isotypic_expansion”) controlling storage/perf trade-offs.

Type:

str

common_irreps#

Irrep identifiers present in both in_rep and out_rep.

Type:

list[tuple]

iso_blocks#

Per-irrep metadata for irreps shared by in_rep/out_rep, keys:

  • out_slice / in_slice: slice selecting the isotypic coordinates in out/in reps.

  • mul_out / mul_in: multiplicities of the irrep in out/in reps.

  • irrep_dim: Dimension of the irreducible representation \(d_k\).

  • endomorphism_basis: basis of \(\operatorname{End}_\mathbb{G}(\hat{\rho}_k)\) with shape \((S_k, d_k, d_k)\).

  • dim_hom_basis: Dimension of the homomorphism between isotypic spaces of type \(k\), i.e. \(\dim(\operatorname{Hom}_\mathbb{G}(\rho_k^{\mathcal{X}}, \rho_k^{\mathcal{Y}})) = m_{out} \cdot m_{in} \cdot S_k\).

  • hom_basis_slice: slice containing the degrees of freedom associated to \(\dim(\operatorname{Hom}_\mathbb{G}(\rho_k^{\mathcal{X}}, \rho_k^{\mathcal{Y}}))\) from a vector of shape \((\dim(\operatorname{Hom}_\mathbb{G}(\rho_{\mathcal{X}}, \rho_{\mathcal{Y}})),)\).

Type:

dict

basis_elements#

Full dense basis stack (dim, out_rep.size, in_rep.size) when basis_expansion="memory_heavy".

Type:

Tensor

basis_norm_sq#

Squared Frobenius norms of basis elements when basis_expansion="memory_heavy".

Type:

Tensor

endo_basis_flat_<irrep_id>

Per-irrep flattened endomorphism bases (S_k, d_k*d_k) when basis_expansion="isotypic_expansion".

Type:

Tensor

endo_basis_norm_sq_<irrep_id>

Per-irrep squared norms of the flattened endomorphism bases with isotypic expansion.

Type:

Tensor

Q_in_inv#

Change-of-basis matrix \(Q_{in}^{-1}\) cached as buffer for isotypic expansion.

Type:

Tensor

Q_out#

Change-of-basis matrix \(Q_{out}\) cached as buffer for isotypic expansion.

Type:

Tensor

Construct the equivariant basis and cache block-wise metadata.

Parameters:

  • in_rep (Representation) – Input representation \(\rho_{\mathcal{X}}\) whose isotypic decomposition has size in_rep.size. The change of basis is applied in-place and used to define the rows of the basis matrices.

  • out_rep (Representation) – Output representation \(\rho_{\mathcal{Y}}\), treated analogously to in_rep with total size out_rep.size.

  • basis_expansion (str) – Strategy for realizing the basis. "memory_heavy" stores the full stack of basis elements with shape (dim(Hom_G(in_rep, out_rep)), out_rep.size, in_rep.size), maximizing speed at the cost of one dense matrix per basis element. "isotypic_expansion" keeps only the tiny irrep endomorphism bases (S_k, d_k, d_k) and reconstructs on the fly—much lighter on memory, moderately slower.

Raises:

AssertionError – If in_rep and out_rep do not belong to the same symmetry group.

property dim: int#

Dimension \(\dim(\operatorname{Hom}_\mathbb{G}(\rho_{\mathcal{X}}, \rho_{\mathcal{Y}}))\).

forward(w_dof)[source]#

Return equivariant linear map from coefficients.

The returned matrix satisfies

\[\rho_{\mathcal{Y}}(g)\mathbf{W} = \mathbf{W}\rho_{\mathcal{X}}(g), \quad \forall g\in\mathbb{G}.\]
Parameters:

w_dof (Tensor) – Basis expansion coefficients of shape (D,) or (..., D), where D = dim(Hom_G(in_rep, out_rep)).

Returns:

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

Return type:

Tensor

Example

>>> W = basis(w_dof)
>>> # equivariance constraint:
>>> # rho_out(g) @ W == W @ rho_in(g) for all g
initialize_params(scheme='kaiming_uniform', return_dense=False, leading_shape=None)[source]#

Sample valid parameters in \(\operatorname{Hom}_\mathbb{G}(\rho_{\mathcal{X}},\rho_{\mathcal{Y}})\).

Parameters:

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

  • return_dense (bool) – If True, return dense weights in the original basis; otherwise return basis expansion coefficients.

  • leading_shape (int | tuple | None) – Optional leading dimensions (e.g., batch size or a tuple of dims). None yields no leading dims. Examples: None → (dim,) / (d_out, d_in); B → (B, dim) / (B, d_out, d_in).

Return type:

Tensor

Shapes:

  • return_dense=False → (*leading_shape, dim(Hom_G(out_rep, in_rep)))

  • return_dense=True → (*leading_shape, out_rep.size, in_rep.size)

Returns:

Initialized parameters with the shapes above.

Return type:

Tensor

Parameters:

Notes

  • If return_dense=False, output lives in coefficient space and can be passed to forward().

  • If return_dense=True, returned dense matrices already satisfy \(\rho_{\mathcal{Y}}(g)\mathbf{W}=\mathbf{W}\rho_{\mathcal{X}}(g)\) for all \(g\).

Example

>>> w_dof = basis.initialize_params("xavier_uniform")
>>> W = basis(w_dof)
>>> W2 = basis.initialize_params("xavier_uniform", return_dense=True)
orthogonal_projection(W)[source]#

Project a dense matrix onto \(\operatorname{Hom}_\mathbb{G}(\rho_{\mathcal{X}}, \rho_{\mathcal{Y}})\).

The projection \(\Pi_{\mathrm{Hom}}(\mathbf{W})\) is orthogonal under the Frobenius inner product and enforces

\[\rho_{\mathcal{Y}}(g)\,\Pi_{\mathrm{Hom}}(\mathbf{W}) = \Pi_{\mathrm{Hom}}(\mathbf{W})\,\rho_{\mathcal{X}}(g), \quad \forall g\in\mathbb{G}.\]
Parameters:

W (Tensor) – Weight matrix of shape (..., out_rep.size, in_rep.size) in the original basis.

Returns:

Projection of W onto the equivariant subspace, matching the input shape.

Return type:

Tensor

Note

The projection is orthogonal with respect to the Frobenius inner product. The selected basis_expansion controls how the projection is computed (speed vs. memory) but not the result.

Example

>>> W_proj = basis.orthogonal_projection(W)
>>> # W_proj satisfies:
>>> # rho_out(g) @ W_proj == W_proj @ rho_in(g) for all g
>>> # and is idempotent under projection:
>>> # basis.orthogonal_projection(W_proj) == W_proj