cuequivariance-torch

Tensor Products

SegmentedPolynomial(polynomial[, ...])	PyTorch module that computes a segmented polynomial.
EquivariantTensorProduct(e, *[, layout, ...])	Equivariant tensor product.
TensorProduct(descriptor, *[, device, ...])	PyTorch module that computes the last operand of the segmented tensor product defined by the descriptor.

Special Cases of Tensor Products

ChannelWiseTensorProduct(irreps_in1, irreps_in2)	Channel-wise tensor product layer.
FullyConnectedTensorProduct(irreps_in1, ...)	Fully connected tensor product layer.
Linear(irreps_in, irreps_out, *[, layout, ...])	A class that represents an equivariant linear layer.
SymmetricContraction(irreps_in, irreps_out, ...)	Accelerated implementation of the symmetric contraction operation introduced in https://arxiv.org/abs/2206.07697.
TransposeIrrepsLayout(irreps, *, source, target)	Transpose the irreps layout of a tensor.
SphericalHarmonics(ls[, normalize, device, ...])	Compute the spherical harmonics of the input vectors as a torch module.

Euclidean Operations

Rotation(irreps, *[, layout, layout_in, ...])	A class that represents a rotation layer for SO3 or O3 representations.
Inversion(irreps, *[, layout, layout_in, ...])	Inversion layer for \(O(3)\) representations.

encode_rotation_angle(angle, ell)	Encode a angle into a tensor of cosines and sines.
vector_to_euler_angles(vector)	Convert a 3D vector to Euler angles.

Extra Modules

layers.BatchNorm(irreps, *[, layout, eps, ...])	Batch normalization for orthonormal representations.
layers.FullyConnectedTensorProductConv(...)	Message passing layer for tensor products in DiffDock-like architectures.