cuequivariance_torch: Executing Equivariant Polynomials in PyTorch

Overview

cuequivariance_torch (imported as cuet) executes cuequivariance polynomials on GPU via PyTorch. It provides:

1. Core primitive: cuet.SegmentedPolynomial — torch.nn.Module with multiple CUDA backends
2. High-level operations (torch.nn.Module): ChannelWiseTensorProduct, FullyConnectedTensorProduct, Linear, SymmetricContraction, SphericalHarmonics, Rotation, Inversion
3. Layers: cuet.layers.BatchNorm, cuet.layers.FullyConnectedTensorProductConv (message passing)
4. Utilities: triangle_attention, triangle_multiplicative_update, attention_pair_bias (AlphaFold2-style)
5. Export support: onnx_custom_translation_table(), register_tensorrt_plugins()

Execution methods