fix(tf): align ProdForceGrad with ProdForce layout

[njzjz-bot]

Summary

• align the public legacy ProdForceGrad raw op with ProdForce's [nframes, 3 * nall] output layout
• keep extended atom slots distinct when callers explicitly provide nall > nloc
• validate the raw op's neighbor indices with one parallel reduction

Reachability and compatibility

TensorFlow's registered training gradient currently supplies nall == nloc, so this is contract alignment rather than a change to the ordinary training path. The nall > nloc regression intentionally exercises the separately registered raw op, whose grad input is indexed using nall and must remain memory-safe for its declared layout.

Validation

• focused ProdForceGrad regression at the branch head
• ruff format .
• changed-file Ruff checks
• clang-format dry run with --Werror
• git diff --check

Fixes #5656

Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh

[coderabbitai]

📝 Walkthrough

Walkthrough

ProdForceGrad now supports nall greater than nloc, validates full-atom gradient and neighbor-index bounds, preserves ghost indices, and adds regression coverage for second-ghost gradients.

Changes

ProdForceGrad ghost-atom handling

Layer / File(s)	Summary
Full-atom gradient contract and indexing source/op/tf/prod_force_grad.cc	The operation reads nall, validates nall >= nloc, requires gradients sized for all atoms, checks neighbor bounds, indexes using nall, and removes modulo folding of ghost indices.
Second-ghost regression coverage source/tests/tf/test_prod_force_grad.py	A TensorFlow test verifies that a second-ghost neighbor uses its own upstream gradient and produces the expected force and network-gradient values.

Estimated code review effort: 3 (Moderate) | ~20 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	The code now accepts 3*nall gradients, handles ghost indices consistently, and adds regression coverage for ghost atoms.
Out of Scope Changes check	✅ Passed	The changes stay focused on the ProdForceGrad boundary fix and its test coverage.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the main change: aligning ProdForceGrad with ProdForce’s layout and gradient contract.

✨ 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.}

[codecov]

Codecov Report

❌ Patch coverage is 57.14286% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 78.96%. Comparing base (cc908a8) to head (3e4d64f).
⚠️ Report is 10 commits behind head on master.

Files with missing lines	Patch %	Lines
source/op/tf/prod_force_grad.cc	57.14%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5848      +/-   ##
==========================================
- Coverage   79.03%   78.96%   -0.08%     
==========================================
  Files        1055     1069      +14     
  Lines      122233   124076    +1843     
  Branches     4401     4522     +121     
==========================================
+ Hits        96607    97975    +1368     
- Misses      24061    24484     +423     
- Partials     1565     1617      +52     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[wanghan-iapcm]

I don't think the condition this PR fixes is reachable, so I'd like the rationale revisited before it merges.

ProdForceGrad exists only as @ops.RegisterGradient("ProdForce") (deepmd/tf/op/_prod_force_grad.py:15), so it is instantiated only when tf.gradients builds a backward pass — i.e. training. It never appears in a frozen graph or the C++/LAMMPS path (grep for it under source/api_cc/ and source/lmp/ returns nothing). And in training nall == nloc always: deepmd/utils/data.py builds the natoms vector as tmp = [natoms, natoms], which is why source/op/tf/descrpt.cc:133 can assert(nloc == nall) on the mesh mode Python actually produces, remapping every neighbor into [0, nloc) at lines 389/396.

So the forward/backward asymmetry here looks intentional rather than an oversight: ProdForce serves both training and inference, where ghost atoms are real, hence its 3 * nall output; ProdForceGrad only ever runs in training, where there are no ghosts. That is consistent with 7faa186 (2018) changing the forward and deliberately leaving the gradient alone.

Concretely, with nall == nloc every part of this change is a no-op: nall*3 == nloc*3, kk*nall*3 == kk*nloc*3, the removed j_idx > nloc fold was unreachable (all j_idx < nall == nloc), and the new nlist(ii) < nall guard can never fire. The new test only passes a natoms=[1, 3, ...] that the data pipeline never produces.

I'm not against aligning the two contracts on paper, but please re-describe the PR as contract alignment rather than a ghost-gradient-loss fix, and drop the added per-backward-step validation (see inline) — I'd rather not pay a per-step cost on a legacy path to guard a condition that cannot occur. If you believe there is a reachable path where ProdForceGrad sees nall > nloc, please point at it and I'll happily re-review on that basis.

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@source/op/tf/prod_force_grad.cc`:
- Around line 89-92: Update the comment near the natoms layout checks in the
ProdForce gradient implementation to remove the claim that the registered
TensorFlow gradient never sees ghosts. Describe these checks as reachable
contract enforcement for cases such as nloc=1 and nall=3, while preserving their
role in validating the layout used by grad().

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro Plus

Run ID: e2055a2c-f81f-4810-81e7-bf698ac4ce50

📥 Commits

Reviewing files that changed from the base of the PR and between cc908a8 and 9a66c05.

📒 Files selected for processing (2)

• source/op/tf/prod_force_grad.cc
• source/tests/tf/test_prod_force_grad.py

[wanghan-iapcm]

Thanks for taking the reframing seriously rather than defending the original description. The "Reachability and compatibility" section now says what I wanted it to say -- that the registered training gradient supplies nall == nloc, so this is contract alignment and not a change to the ordinary training path -- and the title matches. That is the right record for whoever reads this next to 7faa186ed.

Withdrawing half of my own request: you were right to keep the neighbor bound rather than drop it. Once the j_idx % nloc fold is gone, j_idx addresses grad() directly, and ProdForceGrad is separately registered as a public op, so a directly-constructed call with an unbounded nlist would read out of bounds. With the existing if (j_idx < 0) continue; in the inner loop, j_idx is now bounded on both sides. The single parallel reduction costs essentially nothing next to the compute loop, which was my actual concern. Using reduction(|:) instead of reduction(max:) for OpenMP 2.0 compatibility is a good call -- that matters for the MSVC builds.

The comment on the two guards is also what I was after: it tells the next reader the checks describe the layout grad() is indexed with, rather than implying a failure mode that the training path can hit.

On testing: there is no missed regression to point at here, since the condition is unreachable through the registered gradient. The new case exercises the raw op directly with natoms=[1, 3, ...], which the data pipeline never produces, and it does fail before the change because the old op expected a 3 * nloc grad and folded the ghost index. That makes it a sound contract test, which matches how the PR now describes itself.