pooling

class symm_learning.nn.pooling.IrrepSubspaceNormPooling(in_type: FieldType)

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, ...]

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

export() → Module

Exporting to a torch.nn.Module

extra_repr() → str

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

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