Make NNPOps compatible with newer versions of TorchANI

[epretti]

Presently NNPOps will not work with versions of TorchANI newer than 2.2.4, since there were breaking changes introduced after this (see TorchANI readme and migration guide). TorchANI 2.2.4 is now over 2 years old and the latest version is now 2.7.9, so we should update NNPOps accordingly. A few comments and questions:

• Newer versions of TorchANI are only distributed on PyPI, not conda. As such, we'll need to eventually make sure NNPOps can be installed via pip. Since we want to ensure TorchANI is not a required dependency to use NNPOps (it should only need to be there if you actually want to use the NNPOps replacements of TorchANI classes), it should remain possible to distribute NNPOps via conda.
• Do we want to retain any kind of support for the old versions of TorchANI, or only support versions newer than 2.2.4 from now on? Since NNPOps replaces some of TorchANI's classes with its own, it needs to work a bit with its internals, so it could be tricky to support multiple versions with incompatible APIs.
• TorchANI since v2.6 now includes a "cuAEV" extension that is apparently an optimized CUDA implementation of AEVComputer (see here, here, and here). Have we tested this as to its performance / do we understand whether or not this makes NNPOps' TorchANISymmetryFunctions obsolete or not?
• Should Remove TorchANI imports for type annotations #127 be merged and/or should these changes go into new PR? The current issue with one of the CI runs on Remove TorchANI imports for type annotations #127 seems to be related to the installation of an even older version of TorchANI with a bug.
• Do we have any other repositories where the CI uses pip instead of conda? Currently the CI revolves around conda, and we'll want to install TorchANI via pip for testing, which means presumably that we'll want to install its dependencies like PyTorch and the CUDA toolkit via pip. I can work on this but don't want to start from scratch if we already have something like this somewhere else.

[peastman]

Do we want to retain any kind of support for the old versions of TorchANI, or only support versions newer than 2.2.4 from now on?

I think it's fine to support only recent versions.

Have we tested this as to its performance

No, but it would be an interesting comparison to do.

Should #127 be merged and/or should these changes go into new PR?

Whatever is easiest. We need to make TorchANI no longer a hard dependency. That can be done either as a separate PR or as part of the update to newer versions.

Do we have any other repositories where the CI uses pip instead of conda?

Which one you use makes very little difference. In the CI script, just call pip install -r requirements.txt.

[epretti]

Since I keep running into new problems every time I work on this, I wanted to provide a progress update. After getting past the package version issues, I've gotten TorchANIEnergyShifter and TorchANISpeciesConverter to work, but these aren't very interesting parts. After some tracing through to figure out what they renamed all of the parameters to, I've gotten TorchANISymmetryFunctions to return correct answers. However, there are some problems.

• TorchANI's AEVComputer (what NNPOps tries to replace with TorchANISymmetryFunctions) has been moved from nnp.aev_computer to nnp.nnp.aev_computer, where nnp.aev_computer is a property returning the latter and nnp is an ANI model.
  • Assigning a TorchANISymmetryFunctions to nnp.aev_computer doesn't assign to nnp.nnp.aev_computer but creates a hidden attribute that can only be retrieved with __getattr__.
  • Furthermore, the nnp.aev_computer property has a type annotation and is expected to be the TorchANI native AEVComputer and not NNPOps' version. All of this hopelessly confuses PyTorch if you try to save the result as TorchScript, and I don't know how it can be worked around in NNPOps as it is designed now.
• TorchANI added support for its built-in neighbor list to AEVComputer, and might try to call a nonexistent compute_from_neighbors method on our TorchANISymmetryFunctions if it's already done neighbor calculations. I believe we'd have to support this with different C++ and CUDA kernels. I don't know how much this would be worth the effort given that TorchANI now offers an optimized CUDA implementation of AEVComputer natively.

I haven't figured out how to fix _BatchedNN yet, and so on account of those two, OptimizedTorchANI can't be made to work yet.

A more fundamental problem with all of this is that NNPOps currently reaches into the internals of TorchANI and makes a lot of changes that depend on specific implementation details of TorchANI. Unless we have some guarantee that the internal APIs we're touching remain stable, NNPOps could break at any time in the future if a new TorchANI release changes an internal detail we depend on.

[epretti]

If I had to summarize, I'd say that it seems like updating NNPOps to work with new TorchANI versions is difficult (but probably still possible if I keep working on it), and the benefits to using NNPOps with new TorchANI versions are unclear due to apparent improvements in TorchANI itself.

Maybe I can try to run some benchmarks of an old version of TorchANI with and without NNPOps, vs. a new version with their (supposedly) optimized operations, to see where things stand? Unless we know we want to update NNPOps' TorchANI support regardless; then I will just keep working on it.

[peastman]

I don't think it's worth a lot of effort. ANI isn't used nearly as much anymore, and it's not clear whether this still provides a significant benefit. The new release of OpenMM-ML switches to using the TorchANI implementation by default.

[epretti]

The dependency on TorchANI has been removed in #129.