fix(electrostatics): dechain dependent autograd inputs

[zubatyuk]

ALCHEMI Toolkit-Ops Pull Request

Description

Fix the Torch custom-backward fallback for energy-only PME and Ewald outputs. It is used for non-uniform energy cotangents and when the caller requests a differentiable first derivative (create_graph=True). When charges depended on positions, positions depended on cell/strain, or both, the fallback returned connected total gradients for multiple inputs. PyTorch then propagated each returned downstream gradient through its upstream dependency again.

The bug affected first-order position and cell gradients from weighted energy losses, plus second-order derivatives (HVPs and force-loss/double-backward workloads). It covered the shared Ewald/PME energy connector and PME reciprocal cached-first-gradient connector. Uniform first-order cached gradients and direct force/virial output semantics are unchanged.

Type of Change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Performance improvement
• Documentation update
• Refactoring (no functional changes)
• CI/CD or infrastructure change

Related Issues

None.

Changes Made

• Return independent partial gradients by removing downstream-to-upstream paths in charge, position, cell, and alpha fallback routing.
• Apply the fix to the shared Ewald/PME energy connector and PME reciprocal cached-first-gradient connector while preserving uniform cached first gradients and direct outputs.
• Add Torch and JAX regressions for charge-dependent positions, cell-derived positions, weighted backward, HVPs, and double backward.

Testing

• Unit tests pass locally (make pytest)
• Linting passes (make lint)
• New tests added for new functionality meets coverage expectations?

Targeted Torch electrostatics utility and PME suites pass locally.

Checklist

• I have read and understand the Contributing Guidelines
• I have updated the CHANGELOG.md
• I have performed a self-review of my code
• I have added docstrings to new functions/classes
• I have updated the documentation (if applicable)

Additional Notes

None.

Tip

This repository uses Greptile, an AI code review service, to help conduct
pull request reviews. We encourage contributors to read and consider suggestions
made by Greptile, but note that human maintainers will provide the necessary
reviews for merging: Greptile's comments are not a qualitative judgement
of your code, nor is it an indication that the PR will be accepted/rejected.
We encourage the use of emoji reactions to Greptile comments, depending on
their usefulness and accuracy.

[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.

[greptile-apps]

Greptile Summary

This PR fixes double-counted autograd contributions in the Torch custom-backward fallback for energy-only PME and Ewald outputs. The root cause was that returning total connected gradients for all differentiable inputs (positions, charges, cell, alpha) when those inputs share autograd paths caused PyTorch to propagate each returned gradient through its upstream sub-graph a second time.

• Introduces _dechain_connected_input_grads, which processes inputs in downstream-to-upstream order (charges \u2192 positions \u2192 cell \u2192 alpha) and subtracts from each parent\u2019s gradient the VJP contribution already carried by its child, leaving only independent partial gradients for PyTorch to chain through.
• Replaces the multi-branch special-casing in _InjectCachedEvalGradWithFallback and _PMEReciprocalCachedFirstGrad with a single unified recompute + dechain pass; also changes the particle_mesh_ewald fallback closure to call _pme_reciprocal_space_impl directly instead of pme_reciprocal_space, preventing double-entry into the reciprocal cached-first connector.
• Adds thorough regression coverage in both Torch (sibling graphs, q(cell) gradients, HVPs, cell\u2192position\u2192charge chains, weighted backward, double-backward) and JAX (sibling chain decomposition, JIT, HVP vs FD).

Important Files Changed

Filename	Overview
nvalchemiops/torch/interactions/electrostatics/_util.py	Core fix: adds _dechain_connected_input_grads to remove dependent-input paths from connected gradients; simplifies _has_potentially_geometry_dependent_charges but leaves positions as a now-unused parameter in its signature
nvalchemiops/torch/interactions/electrostatics/pme.py	Simplifies _PMEReciprocalCachedFirstGrad.backward using _dechain_connected_input_grads; replaces nested pme_reciprocal_space call in _fallback with direct _pme_reciprocal_space_impl to avoid re-entering the cached-first connector
nvalchemiops/torch/interactions/electrostatics/ewald.py	Documentation updates only; uses existing _has_potentially_geometry_dependent_charges whose behaviour is now subtly more conservative (no longer gates on positions.requires_grad)
test/interactions/electrostatics/test_util.py	Adds four targeted unit tests for _InjectCachedEvalGradWithFallback: sibling position graph, cell-dependent charges, nested cell→positions→charges dechain, and weighted backward dechain; expected values are analytically verified
test/interactions/electrostatics/bindings/torch/test_pme.py	Adds comprehensive PME regression tests: sibling q(R) forces vs FD, q(cell) gradients and HVPs, positions-derived-from-cell, full q(R) cell dependency vs eager oracle, nested cached-first avoidance, and alpha dtype preservation
test/interactions/electrostatics/bindings/torch/test_ewald.py	Adds Ewald q(R) sibling position force vs FD test covering real, reciprocal, and summation variants
test/interactions/electrostatics/bindings/jax/conftest.py	Adds toy_charge_model_jax and qr_manual_chain_jax helpers for q(R) chain-rule JAX tests
test/interactions/electrostatics/bindings/jax/test_ewald.py	Adds JAX Ewald q(R) tests: sibling gradient vs manual chain decomposition, and position HVP vs finite differences
test/interactions/electrostatics/bindings/jax/test_pme.py	Adds JAX PME q(R) parametrised tests for reciprocal-only and full PME: sibling gradient vs manual chain and position HVP vs FD, with and without JIT
CHANGELOG.md	Adds Unreleased section describing the PME/Ewald autograd fix for charge-equilibration terms under create_graph=True

Comments Outside Diff (2)

1. nvalchemiops/torch/interactions/electrostatics/_util.py, line 51-54 (link)

   The positions parameter is now silently unused — the implementation only inspects charges. Keeping it in the signature misleads callers (including the ewald.py call-site) into thinking positions.requires_grad still gates the result. Since the function is private and the only remaining caller passes two arguments, it would be cleaner to drop the dead parameter and update the ewald.py call-site accordingly.

2. nvalchemiops/torch/interactions/electrostatics/ewald.py, line 855 (link)

   Fast-path now slightly more conservative after _has_potentially_geometry_dependent_charges change. The updated implementation returns True whenever charges.requires_grad and charges.grad_fn is not None, regardless of positions.requires_grad. Any call where positions.requires_grad=False but charges is a non-leaf depending only on an independent parameter (e.g., charges = f(theta)) will now skip the cached fast-path and fall through to the eager connected forward — a performance regression for that case, though not a correctness issue. Consider documenting this intentional conservatism, or restoring the positions.requires_grad check in the outer condition so the fast-path guard stays accurate.

   Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

_{Reviews (1): Last reviewed commit: "fix(electrostatics): dechain dependent a..." | Re-trigger Greptile}