NVIDIA / NVIDIA/cuEquivariance
`SphericalHarmonics` naive backend does not support `torch.compile`
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 433
- Forks
- 42
- PR merge metrics
- No merged PRs in 30d
Description
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
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with descriptors.spherical_harmonics and cuet.SegmentedPolynomial(..., method="naive"), then run the provided CPU torch.compile reproduction with PyTorch 2.8. Trace where the naive spherical harmonics path invokes Subscripts and review the custom_op integration there. Done means the reproduction succeeds with fullgraph=True without requiring CUDA.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- backend, machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 58/100