pooling
- class symm_learning.nn.pooling.IrrepSubspaceNormPooling(in_type: FieldType)[source]
Module that outputs the norm of the features in each G-irreducible subspace of the input tensor.
- Parameters:
in_type – Input FieldType. The dimension of the output tensors will be equal to the number of irreps in this type
- evaluate_output_shape(input_shape: tuple[int, ...]) tuple[int, ...] [source]
Compute the shape the output tensor which would be generated by this module when a tensor with shape
input_shape
is provided as input.- Parameters:
input_shape (tuple) – shape of the input tensor
- Returns:
shape of the output tensor
- extra_repr() str [source]
Return the extra representation of the module.
To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.
- forward(x: GeometricTensor) GeometricTensor [source]
Computes the norm of each G-irreducible subspace of the input GeometricTensor.
The input_type representation in the spectral basis is composed of direct sum of N irreducible representations. This function computes the norms of the vectors on each G-irreducible subspace associated with each irrep.
- Parameters:
x – Input GeometricTensor.
- Returns:
G-Invariant tensor of shape (…, N) where N is the number of irreps in the input type.
- Return type:
GeometricTensor