fix(tf): handle multiframe DeepSpin extension - #5851
Conversation
Build TensorFlow DeepSpin virtual-atom coordinates for every frame and use explicit frame strides when copying physical, magnetic, and atomic outputs back from extended buffers. Derive atomic-energy frame width from the TensorFlow output so physical-only energies do not interleave frames in the extended atom layout. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5851 +/- ##
==========================================
+ Coverage 78.58% 79.22% +0.64%
==========================================
Files 1050 1072 +22
Lines 120637 125176 +4539
Branches 4356 4552 +196
==========================================
+ Hits 94801 99171 +4370
- Misses 24278 24374 +96
- Partials 1558 1631 +73 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
Possible reviewers based on changed lines, exact file history, and exact-file review history:
No review request was made automatically. Coding agent: Codex |
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
The core arithmetic here is right, and I checked the central assumption rather than taking it on faith: deepmd/tf/model/ener.py reshapes o_atom_energy to the real-type atom count and the spin fitting emits energies for physical atoms only, so deriving the source stride from output_ae.NumElements() / nframes and leaving the virtual slots zero is correct - and it removes an out-of-bounds read on oae that was there even at nframes == 1. The common.cc change is a provable no-op for the other nine callers of select_real_atoms_coord, since each passes dcoord_ sized exactly nframes * nall * 3.
The regression test is a real one. On master extend_nlist resizes extend_dcoord to extend_nall * 3 and reads only dcoord_[ii * 3 + jj], so a two-frame call returns a one-frame buffer, session_input_tensors infers nframes == 1, and ASSERT_EQ(energy.size(), 2U) fails. The EXPECT_GT(fabs(energy_first - energy_second), EPSILON) guard also stops a duplicated-frame implementation from passing.
One gap inline. Two smaller notes, neither blocking:
-
Both new tests run with
nghost == 0and a model withdim_aparam == 0, so thenew_idx + (ii < nloc ? nloc : nghost)ghost-virtual branch you rewrote is never taken with ghosts present, and thedaparam > 0branch ofselect_real_atoms_coord- the exact case the new comment cites as motivation ("aparam keeps the caller's original atom stride ... DeepSpin virtual atoms") - is never entered. A LAMMPS-style case withnghost > 0would cover the more fragile of the two. -
The new size checks in
extend()andextend_nlist()useassert, which is compiled out under-DNDEBUG, whilerun_modelin the same diff throwsdeepmd::deepmd_exceptionfor its new checks. Mixed, but the asserts guard caller-supplieddcoord_/dspin_sizes, which is where athrowearns its keep.
Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
📝 WalkthroughWalkthroughDeepSpin TensorFlow execution now supports multiple coordinate and spin frames. Extension helpers build frame-indexed virtual atoms, execution paths remap outputs per frame, and tests compare batched results with independent computations. ChangesMulti-frame DeepSpin execution
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant DeepSpinTF
participant ExtensionHelpers
participant TensorFlow
participant OutputRemapping
DeepSpinTF->>ExtensionHelpers: Extend coordinates and spins for nframes
ExtensionHelpers-->>TensorFlow: Provide frame-indexed virtual-atom inputs
TensorFlow-->>OutputRemapping: Return batched model outputs
OutputRemapping-->>DeepSpinTF: Copy real-atom outputs for each frame
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
source/api_cc/src/common.cc (2)
187-194: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winThree new size-consistency checks use
assert, which is inactive in release builds. All three validate externally-influenced coordinate/spin buffer sizes against the derivednframes, butassertcompiles out underNDEBUG, so a caller-supplied size mismatch silently truncates the derived stride instead of failing clearly. The same PR already usesdeepmd::deepmd_exceptionfor an analogous check inrun_model(DeepSpinTF.cclines 179-188); mirroring that pattern here keeps the checks active in release builds.
source/api_cc/src/common.cc#L187-L194: replace theassertat line 192 with a throwndeepmd::deepmd_exceptionwhendcoord_.size()is not evenly divisible bynframes*3.source/api_cc/src/DeepSpinTF.cc#L1074-L1077: replace theassertvalidatingdcoord/spinsizes inextend()with a thrown exception.source/api_cc/src/DeepSpinTF.cc#L1296-L1297: replace theassertvalidatingdcoord_/dspin_sizes inextend_nlist()with a thrown exception.🤖 Prompt for 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. In `@source/api_cc/src/common.cc` around lines 187 - 194, Replace the size-validation asserts with active deepmd::deepmd_exception checks in common.cc lines 187-194, DeepSpinTF.cc lines 1074-1077, and DeepSpinTF.cc lines 1296-1297. Validate divisibility and coordinate/spin size consistency before deriving or using strides, following the existing run_model exception pattern; preserve the current valid-input behavior.
192-192: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winConsider an exception instead of
assertfor this size check.This validates that
dcoord_'s size is evenly divisible bynframes*3. UnderNDEBUG,assertis compiled out, so a mismatched caller-supplied buffer would silently truncatecoord_nallinstead of failing clearly. See the analogous exception-based check added inDeepSpinTF.cc'srun_model(same PR, lines 179-188) for a pattern that stays active in release builds.🤖 Prompt for 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. In `@source/api_cc/src/common.cc` at line 192, Replace the assert in the surrounding coordinate-validation logic with an exception-based size check that remains active under NDEBUG. Validate that dcoord_.size() exactly equals nframes multiplied by coord_nall and 3, and throw a clear error on mismatch before any truncating computation proceeds.
🤖 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.
Nitpick comments:
In `@source/api_cc/src/common.cc`:
- Around line 187-194: Replace the size-validation asserts with active
deepmd::deepmd_exception checks in common.cc lines 187-194, DeepSpinTF.cc lines
1074-1077, and DeepSpinTF.cc lines 1296-1297. Validate divisibility and
coordinate/spin size consistency before deriving or using strides, following the
existing run_model exception pattern; preserve the current valid-input behavior.
- Line 192: Replace the assert in the surrounding coordinate-validation logic
with an exception-based size check that remains active under NDEBUG. Validate
that dcoord_.size() exactly equals nframes multiplied by coord_nall and 3, and
throw a clear error on mismatch before any truncating computation proceeds.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: d1533e3d-e80c-4945-8cd1-1c967e157145
📒 Files selected for processing (4)
source/api_cc/include/DeepSpinTF.hsource/api_cc/src/DeepSpinTF.ccsource/api_cc/src/common.ccsource/api_cc/tests/test_deeppot_tf_spin.cc
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Addressed, and the test half is the part I like most.
The extension now trims each frame's atomic outputs to the real atoms with ff * extend_nall + ii as the source index, so the two overloads in this file finally agree on what they hand back. What makes it convincing rather than merely plausible is that the assertions were tightened at the same time: on master this path had // EXPECT_EQ(atom_ener.size(), natoms); commented out and EXPECT_EQ(atom_vir.size(), (natoms + 2) * 9); asserting the oversized length, while the neighbor-list overload a few hundred lines down already asserted the correct natoms and natoms * 9. In other words the suite had been adjusted to match the broken output, which is exactly why nothing caught this. Both are now exact at all three call sites, so the old behaviour cannot come back quietly.
I did not build TensorFlow C++ locally, so on the question of whether the new assertions fail unpatched I am reasoning from the stride rather than from a run: pre-fix datom_energy_ keeps the extend_nall = nloc + nloc_spin stride, so its length is natoms + 2 and the re-enabled check cannot pass. Worth stating plainly rather than implying I executed it.
On the red checks: the two failing Test C++ jobs are the Paddle inference library download (file DOWNLOAD cannot compute hash on failed download), not this change. Test C++ (true, false, false, true) passed, and it did install DeepSpinTF.h and run all three test binaries with runUnitTest_cc green, so despite the note about the local build not reaching compilation, this code is compiled and exercised in CI.
Approving.
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Correcting myself on one point in my approval above, since it was wrong in a way you might have relied on.
I wrote that Test C++ (true, false, false, true) "did install DeepSpinTF.h and run all three test binaries ... so this code is compiled and exercised in CI". That is not right. The header install is unconditional, so it proves nothing. That job configures with:
cmake -D ENABLE_TENSORFLOW=FALSE -D ENABLE_PYTORCH=FALSE -D ENABLE_PADDLE=TRUE ...
and source/api_cc/CMakeLists.txt removes src/DeepSpinTF.cc from the main LIB_SRC list, compiling it only into deepmd_backend_tf under if(ENABLE_TENSORFLOW). With TensorFlow off, the file this PR changes is never compiled, so the green runUnitTest_cc says nothing about it.
The rest of the picture: of the four Test C++ jobs, one failed on the Paddle inference library download and two were cancelled by fail-fast; the only one that completed is the Paddle-only configuration above. I also said "the two failing jobs" when it was one failure plus two cancellations.
So the accurate statement is that no CI job has compiled or run this change at this head, which combined with the local build not reaching compilation means it currently rests entirely on code review.
I am leaving the approval in place: the fix does what I asked, the two overloads now agree, and the tightened assertions (atom_ener.size() == natoms re-enabled, atom_vir.size() == natoms * 9 replacing the oversized (natoms + 2) * 9) cannot pass on the old stride, so I am confident in it by inspection. But it would be worth getting a TensorFlow-enabled Test C++ job green before merging rather than treating the current checks as evidence.
ca7f708
3 participants
Closes #5660.
Summary
aparamstrides during NULL-atom selectiondoubleandfloatWhy existing tests missed this
All TensorFlow DeepSpin C++ tests used scalar-energy, single-frame calls. The extension helpers therefore only needed frame-zero buffers, and frame-zero copyback filled every asserted output. No test exercised the documented vector-energy overload with
nframes > 1, so undersized extension buffers, incorrect inferred frame counts, and missing output frame offsets remained invisible.The new tests perturb both coordinates and spins in frame 2, evaluate each frame separately, and compare every batched slice against its one-frame reference. They also assert the frame references differ, preventing duplicated-frame behavior from passing.
Validation
ruff format .ruff check .clang-format --dry-run --Werroron all changed C++ filesdeepmd_op,deepmd_backend_tf, andrunUnitTests_ccTestInferDeepSpinandTestInferDeepSpinNopbctyped tests passedCoding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh
Summary by CodeRabbit
New Features
Bug Fixes