fix(pt): clone inference coords before enabling grad

[njzjz-bot]

Summary

• clone PyTorch atomic-model coordinates before enabling gradients, so LAMMPS-provided view tensors are not modified in-place
• pass the cloned force-coordinate tensor through private atomic-model metadata for derivative generation
• add regression tests for leaf-view coordinate inputs in atomic and lower model paths

Tests

• uvx pre-commit run --files deepmd/pt/model/atomic_model/base_atomic_model.py deepmd/pt/model/atomic_model/dp_atomic_model.py deepmd/pt/model/atomic_model/linear_atomic_model.py deepmd/pt/model/atomic_model/pairtab_atomic_model.py deepmd/pt/model/model/make_model.py source/tests/pt/model/test_dp_atomic_model.py source/tests/pt/model/test_dp_model.py
• PYTHONPATH=$PWD uv run --index-strategy unsafe-best-match --with pytest --with numpy --with scipy --with pyyaml --with dargs --with typing_extensions --with h5py --with wcmatch --with packaging --with ml_dtypes --with mendeleev --with array-api-compat --with lmdb --with msgpack --with torch==2.5.1+cpu --extra-index-url https://download.pytorch.org/whl/cpu --no-project pytest source/tests/pt/model/test_dp_model.py::TestDPModel::test_forward_lower_accepts_leaf_view_input source/tests/pt/model/test_dp_atomic_model.py::TestDPAtomicModel::test_forward_common_atomic_accepts_leaf_view_input -q

Closes #5165

Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.5)

Summary by CodeRabbit

Release Notes

• Bug Fixes

  • Improved gradient computation for coordinate and charge derivatives across atomic model implementations
  • Enhanced tracking of coordinate tensors during force calculation workflows
  • Refined internal data structure filtering to prevent unintended operations on private entries

• Tests

  • Added test coverage for edge cases involving leaf tensor view inputs in model forward operations

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: b4b1acda-2405-4591-ae3e-0331015cb8ab

📥 Commits

Reviewing files that changed from the base of the PR and between 93f5580 and 3877080.

📒 Files selected for processing (7)

• deepmd/pt/model/atomic_model/base_atomic_model.py
• deepmd/pt/model/atomic_model/dp_atomic_model.py
• deepmd/pt/model/atomic_model/linear_atomic_model.py
• deepmd/pt/model/atomic_model/pairtab_atomic_model.py
• deepmd/pt/model/model/make_model.py
• source/tests/pt/model/test_dp_atomic_model.py
• source/tests/pt/model/test_dp_model.py

---

📝 Walkthrough

Walkthrough

This PR resolves PyTorch RuntimeError occurring when in-place requires_grad_ operations are applied to leaf variable views. The fix uses detach-clone to create independent gradient-tracked tensors, tracks these via _force_coord keys through atomic models, updates masking to preserve private entries, and routes the tracked coordinate through model output assembly.

Changes

Coordinate Gradient and Force Tracking in Atomic Models

Layer / File(s)	Summary
Atomic model coordinate gradient handling deepmd/pt/model/atomic_model/dp_atomic_model.py, deepmd/pt/model/atomic_model/linear_atomic_model.py, deepmd/pt/model/atomic_model/pairtab_atomic_model.py	When do_grad_r() or do_grad_c() is enabled, extended_coord is now detached and cloned with requires_grad_(True) instead of modifying in-place. The gradient-enabled tensor is stored under "_force_coord" key in the output dict.
Base atomic model masking of private keys deepmd/pt/model/atomic_model/base_atomic_model.py	The output post-processing loop in forward_common_atomic now iterates over a key snapshot and skips entries whose names start with "_", protecting internal values like _force_coord from being reshaped or masked.
Model output assembly with force coordinate routing deepmd/pt/model/model/make_model.py	forward_common_lower extracts optional _force_coord from atomic_ret (falling back to cc_ext) and passes it to fit_output_to_model_output, while extended_coord is preserved in the computation graph by removing it from cleanup.
Integration tests for leaf view coordinate inputs source/tests/pt/model/test_dp_atomic_model.py, source/tests/pt/model/test_dp_model.py	New tests confirm that DPAtomicModel.forward_common_atomic and EnergyModel.forward_lower accept coordinate tensors created as leaf variables with gradients enabled and reshaped via .view(), verifying end-to-end operation without in-place mutation errors.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Suggested labels

bug, Python

Suggested reviewers

• njzjz
• iProzd

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and specifically describes the main change: cloning inference coordinates before enabling gradient tracking.
Linked Issues check	✅ Passed	Code changes prevent in-place gradient operations on LAMMPS view tensors by cloning coordinates with detach(), storing references via '_force_coord', and include regression tests validating leaf-view tensor acceptance [#5165].
Out of Scope Changes check	✅ Passed	All changes focus on preventing in-place gradient operations and tracking force coordinates. No unrelated functionality or scope creep detected across modified atomic models and supporting infrastructure.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}