cuequivariance-torch

Tensor Products

SegmentedPolynomial

PyTorch module that computes a segmented polynomial.

Special Cases of Tensor Products

ChannelWiseTensorProduct	Channel-wise tensor product layer.
FullyConnectedTensorProduct	Fully connected tensor product layer.
Linear	A class that represents an equivariant linear layer.
SymmetricContraction	Accelerated implementation of the symmetric contraction operation introduced in https://arxiv.org/abs/2206.07697.
TransposeIrrepsLayout	Transpose the irreps layout of a tensor.
SphericalHarmonics	Compute the spherical harmonics of the input vectors as a torch module.

Euclidean Operations

Rotation	A class that represents a rotation layer for SO3 or O3 representations.
Inversion	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.

Triangle

triangle_attention(q, k, v, bias[, mask, ...])	Triangle Attention
triangle_multiplicative_update(x[, ...])	Apply triangle multiplicative update operation.
attention_pair_bias(s, q, k, v, z, mask, ...)	Compute attention with pairwise bias for diffusion models.

Extra Modules

layers.BatchNorm	Batch normalization for orthonormal representations.
layers.FullyConnectedTensorProductConv	Message passing layer for tensor products in DiffDock-like architectures.