
from cuequivariance388
PyTorch GPU-accelerated equivariant tensor primitives and layers (SegmentedPolynomial, tensor products, spherical harmonics, equivariant linear layers) for buil
Provides GPU-accelerated implementations of equivariant polynomial/tensor operations for PyTorch (imported as cuet). It exposes a SegmentedPolynomial primitive with multiple CUDA backends and a set of high-level torch.nn.Module operations (ChannelWiseTensorProduct, FullyConnectedTensorProduct, Linear, SymmetricContraction, SphericalHarmonics, Rotation/Inversion) plus equivariant layers (BatchNorm, FullyConnectedTensorProductConv). Typical users build equivariant GNNs, message-passing layers, and physics-aware models (e.g., DiffDock, MACE, NEQUIP) that require SO(3)/O(3)-equivariant math on GPU.
Use this skill when you need highly-efficient, correct equivariant tensor contractions and layers in PyTorch, especially for GPU-accelerated workloads: converting descriptor polynomials into torch modules, implementing message-passing in equivariant GNNs, or exporting models to ONNX/TensorRT. Choose this over handwritten einsum code when you need performance and multiple backend strategies (naive, uniform_1d, fused_tp, indexed_linear).
Likely compatible with developer/engineering agent workflows and local execution CLIs that can run PyTorch code (Copilot/Codex/Gemini CLI style developer agents).
NVIDIA cuEquivariance Torch skill providing comprehensive API documentation for equivariant tensor operations in PyTorch. No bundled scripts — purely a reference/documentation skill with well-structured code examples covering SegmentedPolynomial, tensor products, spherical harmonics, and GNN layers. Clean security profile with no concerning patterns. Very niche audience (physics-ML researchers with CUDA GPUs).
Documentation-only skill from NVIDIA's official cuEquivariance repo. High technical quality but extremely niche — serves as API reference rather than an agent-automated workflow skill. The SKILL.md is well-structured with clear code examples and execution method comparison table.