Fix: resolve TorchScript JIT stability and PyTorch 2.2 compatibility bugs

[changzhiai]

Summary

This Pull Request addresses two critical issues in the cuequivariance-torch library that were preventing successful training on certain PyTorch versions and blocking TorchScript export for deployment (e.g., for use in LAMMPS).

---

1. TorchScript JIT Stability & Type Inference

• Problem: In tp_channel_wise.py, the JIT compiler failed to infer types for empty dictionaries. Additionally, a variable shadowing issue existed where indices_out was being used as both an input argument (Tensor) and an internal dictionary, which is not permitted in TorchScript.
• Fix:
  • Renamed the input argument to indices_0 to ensure stable typing.
  • Added explicit type hints (e.g., Dict[int, torch.Tensor]) to internal dictionaries to assist the JIT compiler.

2. PyTorch 2.2+ Compatibility & Scripting Support

• Problem: In segmented_polynomial.py, a call to the private function torch.fx._symbolic_trace.is_fx_symbolic_tracing() caused an AttributeError in PyTorch 2.2.x.
• Fix:
  • Used getattr with the built-in bool function as a default fallback. This safely returns False if the function is missing and is fully compatible with torch.jit.script.

---

Files Modified:

• cuequivariance_torch/operations/tp_channel_wise.py
• cuequivariance_torch/primitives/segmented_polynomial.py

Verification:

• Verified that models can be successfully exported to TorchScript using torch.jit.script for LAMMPS deployment.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[phiandark]

Hello @changzhiai , thank you very much for the PR, all changes look good to me.
Could you please "signoff" your commits as mentioned in the contributing guide?
After that I should be able to approve this.
Thanks!

[changzhiai]

Hi @phiandark, thanks a lot for the guide. Now I already "signoff" the commits.

[phiandark]

Hi @changzhiai , thank you for signing off.

Our CI caught another issue: this changes the named argument of ChannelWise and that can break downstream use, which might be bad for existing implementations. If it's just a matter of duplicate names, would you mind leaving the named argument as indices_out and change the internal version to something else instead?

Also, the pre-commit style check fails: can you make sure to install and run the pre-commit hooks before your next push. Then I can re-approve and merge.
Thank you!

[changzhiai]

Hi @phiandark, thank you so much for your suggestion. I have updated the PR with the following fixes:

1. Backward Compatibility: Reverted the indices_out argument rename in ChannelWiseTensorProduct and handled the naming conflict internally to prevent breaking downstream usage.
2. Style Check: ran the pre-commit hooks (ruff and ruff format). All style checks are now passing locally.
3. PyTorch Compatibility: Implemented a more robust getattr check for is_fx_symbolic_tracing in SegmentedPolynomial. This should resolve the documentation build failure by ensuring compatibility with older PyTorch versions (like 2.2).

Thank you!