Skip to content

Remove TorchANI replacement modules and dependency#129

Merged
epretti merged 16 commits intoopenmm:masterfrom
epretti:remove-torchani
Feb 4, 2026
Merged

Remove TorchANI replacement modules and dependency#129
epretti merged 16 commits intoopenmm:masterfrom
epretti:remove-torchani

Conversation

@epretti
Copy link
Copy Markdown
Member

@epretti epretti commented Jan 16, 2026

Makes NNPOps no longer dependent on TorchANI, and removes parts meant to replace components of TorchANI. For posterity, we decided to do this because keeping the TorchANI dependency and updating it for newer TorchANI versions would be difficult, make NNPOps reliant on internal implementation details of TorchANI that could change in the future, and have unclear performance benefits.

Also:

  • Keeps the ANI symmetry functions accessible from Python through a PyTorch module ANISymmetryFunctions that does not depend on TorchANI
  • Makes the tests for the ANI symmetry functions not depend on TorchANI, and moves the TorchANI-dependent script used to generate the test cases into devtools
  • For ANI symmetry functions, prevents a race condition by fixing some calls to cudaMemcpyAsync and cudaMemsetAsync that would be enqueued on the default CUDA stream even if a different stream was selected for the kernels
  • For ANI symmetry functions and CFConv, fixes checking the results of cudaPointerGetAttributes due to a change in behavior when passed host pointers in CUDA 11 and above that could lead to a race condition
  • Fixes an issue in getNeighborPairs caused by passing a NumPy integer type to PyTorch

I made some changes to get the CI to work on the latest version of PyTorch but did not try to update the .github/workflows/self-hosted-gpu-test.yml; I don't know whether or not that's currently in use.

@peastman
Copy link
Copy Markdown
Member

peastman commented Feb 4, 2026

This all looks good. Go ahead and merge if it's ready.

@epretti epretti merged commit 45feca6 into openmm:master Feb 4, 2026
3 checks passed
@epretti epretti deleted the remove-torchani branch February 4, 2026 19:12
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants