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}})\).

Let \((\mathcal{X}, \rho_{\mathcal{X}})\) and \((\mathcal{Y}, \rho_{\mathcal{Y}})\) be two \(\mathbb{G}\)-symmetric vector spaces. After rewriting both representations in their canonical isotypic decomposition, every equivariant linear map \(\mathbf{W}\in\operatorname{Hom}_\mathbb{G}(\rho_{\mathcal{X}}, \rho_{\mathcal{Y}})\) admits a block decomposition of the form

\[\tilde{\mathbf{W}} = \mathbf{Q}_{\mathcal{Y}}^T \mathbf{W} \mathbf{Q}_{\mathcal{X}} = \bigoplus_{k\in\mathcal{K}_{\mathcal{X}\mathcal{Y}}} \tilde{\mathbf{W}}^{(k)},\]

where \(\mathcal{K}_{\mathcal{X}\mathcal{Y}}\) contains the irrep types shared by \(\rho_{\mathcal{X}}\) and \(\rho_{\mathcal{Y}}\). This is the structural statement summarized in Proposition I.13 of main_vk.pdf and exploited throughout this class.

For each shared irrep type \(k\), let \(m_k^{\mathcal{X}}\) and \(m_k^{\mathcal{Y}}\) denote the multiplicities of \(\hat{\rho}_k\) in \(\rho_{\mathcal{X}}\) and \(\rho_{\mathcal{Y}}\), and let \(\{\mathbf{\Psi}^{(k)}_s\}_{s=1}^{S_k}\) be a basis of \(\operatorname{End}_\mathbb{G}(\hat{\rho}_k)\). Then each isotypic block further decomposes into \(m_k^{\mathcal{Y}} \times m_k^{\mathcal{X}}\) endomorphism blocks:

\[\tilde{\mathbf{W}}^{(k)}_{j,i} = \sum_{s=1}^{S_k}\theta_{j,i,s}^{(k)}\,\mathbf{\Psi}^{(k)}_s, \qquad \tilde{\mathbf{W}}^{(k)} = \sum_{s=1}^{S_k}\mathbf{\Theta}^{(k)}_s \otimes \mathbf{\Psi}^{(k)}_s,\]

with \(\mathbf{\Theta}^{(k)}_s\in\mathbb{R}^{m_k^{\mathcal{Y}}\times m_k^{\mathcal{X}}}\). Equation (40) in main_vk.pdf describes the possible endomorphism bases for real, complex, and quaternionic irrep types, so \(S_k = \dim(\operatorname{End}_\mathbb{G}(\hat{\rho}_k))\) is typically \(1\), \(2\), or \(4\).

Consequently, the dimension of the equivariant subspace is

\[\dim\!\left(\operatorname{Hom}_\mathbb{G}(\rho_{\mathcal{X}}, \rho_{\mathcal{Y}})\right) = \sum_{k\in\mathcal{K}_{\mathcal{X}\mathcal{Y}}} m_k^{\mathcal{Y}} m_k^{\mathcal{X}} S_k.\]

This module turns that structure into a practical parameterization. It provides:

Synthesis of equivariant maps from basis coefficients.
Orthogonal projection of dense matrices onto the equivariant subspace.
Block-wise metadata exposing how the degrees of freedom are partitioned across shared irrep types.

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.

The tutorial Leveraging the structure of Equivariant Linear maps summarizes this decomposition and shows how it matches the implementation in GroupHomomorphismBasis.

Examples

Build and synthesize an equivariant matrix with forward().

>>> from escnn.group import CyclicGroup
>>> G = CyclicGroup(4)
>>> rep_x = direct_sum([G.regular_representation] * 2)
>>> rep_y = direct_sum([G.regular_representation] * 3)
>>> basis = GroupHomomorphismBasis(rep_x, rep_y, basis_expansion="isotypic_expansion")
>>> irrep_id = next(k for k, meta in basis.iso_blocks.items() if meta["dim_endo_basis"] > 1)
>>> block = basis.iso_blocks[irrep_id]
>>> theta_k = torch.randn(block["mul_out"] * block["mul_in"], block["dim_endo_basis"])
>>> w_dof = torch.zeros(basis.dim)
>>> w_dof[block["hom_basis_slice"]] = theta_k.reshape(-1)
>>> W = basis(w_dof)
>>> # theta_k stores the coefficients of the block associated with irrep_id
>>> # W satisfies rho_y(g) @ W == W @ rho_x(g) for all g in G

Project an unconstrained matrix with orthogonal_projection().

>>> W0 = torch.randn(rep_y.size, rep_x.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 \(m_k^{\mathcal{Y}}\) and \(m_k^{\mathcal{X}}\).
irrep_dim: Dimension of the irreducible representation \(d_k\).
dim_endo_basis: \(S_k = \dim(\operatorname{End}_\mathbb{G}(\hat{\rho}_k))\).
dim_hom_basis: Dimension of the homomorphism between isotypic spaces of type \(k\), i.e. \(m_k^{\mathcal{Y}} m_k^{\mathcal{X}} S_k\).
hom_basis_slice: slice containing the degrees of freedom associated to \(\operatorname{Hom}_\mathbb{G}(\rho_{\mathcal{X}}^{(k)}, \rho_{\mathcal{Y}}^{(k)})\) inside the flattened coefficient vector. The exact flat ordering is backend-dependent, but the slice is always compatible with forward() and projection_coefficients() of the same instance.

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#

Cached matrix \(\mathbf{Q}_{\mathcal{X}}^T\) mapping the original input coordinates to isotypic coordinates when basis_expansion="isotypic_expansion".

Type:: Tensor

Q_out#

Cached matrix \(\mathbf{Q}_{\mathcal{Y}}\) mapping from isotypic coordinates back to the original output basis when basis_expansion="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. Internally, the representation is rewritten in canonical isotypic ordering before any basis metadata is constructed.
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 small irrep endomorphism bases (S_k, d_k, d_k) and reconstructs the dense matrix on the fly, which is much lighter on memory and exposes the isotypic block structure more directly.

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 a dense equivariant linear map from flattened coefficients.

The returned matrix satisfies

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

In the isotypic basis, the map is assembled block by block as

\[\tilde{\mathbf{W}}^{(k)}_{j,i} = \sum_{s=1}^{S_k}\theta_{j,i,s}^{(k)}\,\mathbf{\Psi}^{(k)}_s,\]

and then returned to the original coordinates through

\[\mathbf{W} = \mathbf{Q}_{\mathcal{Y}} \tilde{\mathbf{W}} \mathbf{Q}_{\mathcal{X}}^T.\]

Parameters:: w_dof (Tensor) – Basis expansion coefficients of shape (D,) or (..., D), where D = dim(Hom_G(in_rep, out_rep)). The flat layout is the one expected by the selected basis_expansion backend, and is always compatible with projection_coefficients() from the same instance.
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:

scheme (str)
return_dense (bool)
leading_shape (int | tuple | None)

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

projection_coefficients(W)[source]#

Return homomorphism-basis coefficients of the orthogonal projection of W.

This method first computes the Frobenius-orthogonal projection

\[\Pi_{\mathrm{Hom}_{\mathbb{G}}}(\mathbf{W}) \in \operatorname{Hom}_{\mathbb{G}}(\rho_{\mathcal{X}}, \rho_{\mathcal{Y}}),\]

and then returns the flattened coefficients of that projection in the basis used by this GroupHomomorphismBasis instance. The resulting vector can be passed back to forward() to reconstruct the dense projected matrix exactly up to numerical precision.

Parameters:: W (Tensor) – Dense matrix or batch of matrices of shape (..., out_rep.size, in_rep.size).
Returns:: Flattened coefficient vector(s) of shape (..., dim).
Return type:: Tensor

property tensor_cache: IsotypicTensorCache#: Structured tensor cache for the isotypic-expansion linalg kernels.