impose_linear_equivariance#
- impose_linear_equivariance(lin, in_rep, out_rep, basis_expansion_scheme='isotypic_expansion')[source]#
Impose equivariance constraints on a given torch.nn.Linear layer using torch parametrizations.
Impose via torch parametrizations (hard constraints on trainable parameters ) that the weight matrix of the given linear layer commutes with the group actions of the input and output representations. That is:
\[\rho_{\text{out}}(g) W = W \rho_{\text{in}}(g) \quad \forall g \in G\]If the layer has a bias term, it is constrained to be invariant:
\[\rho_{\text{out}}(g) b = b \quad \forall g \in G\]- Parameters:
lin (
Module) – The linear layer to impose equivariance on. Must have ‘weight’ and optionally ‘bias’ attributes.in_rep (
Representation) – The input representation \(\rho_{\text{in}}\) of the layer.out_rep (
Representation) – The output representation \(\rho_{\text{out}}\) of the layer.basis_expansion_scheme (str) – Basis expansion strategy for the commuting constraint (
"memory_heavy"or"isotypic_expansion").
- Return type: