Describe the bug
The descriptors.spherical_harmonics API causes torch.compile to fail. This is critical because it is automatically used as a fallback in certain cases, meaning it is strictly broken on the CI where cuda is unavailable
To Reproduce
import torch
import cuequivariance as cue
import cuequivariance_torch as cuet
from cuequivariance import descriptors
e = descriptors.spherical_harmonics(cue.SO3(1), [0, 1, 2])
f = cuet.SegmentedPolynomial(e.polynomial, method="naive").to("cpu")
x = torch.randn(4, 3)
# Fails with:
# Unsupported method call: Dynamo does not know how to trace
# method `__add__` of class `Subscripts`.
torch.compile(f, fullgraph=True)([x])
Expected behavior
The spherical_harmonics method should use a custom_op to be compatible with torch.compile.
GPU HW/SW(please complete the following information):
- CUDA toolkit version: Repro does not require CUDA
- torch or ngc docker version: 2.8
- Driver version: Repro does not require CUDA
- full name of GPU: Repro does not require CUDA
Describe the bug
The
descriptors.spherical_harmonicsAPI causestorch.compileto fail. This is critical because it is automatically used as a fallback in certain cases, meaning it is strictly broken on the CI where cuda is unavailableTo Reproduce
Expected behavior
The
spherical_harmonicsmethod should use acustom_opto be compatible withtorch.compile.GPU HW/SW(please complete the following information):