NVIDIA / NVIDIA/cuEquivariance

`SphericalHarmonics` naive backend does not support `torch.compile`

Open
#265 1 comment 0 reactions 0 assignees View on GitHub

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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.