fix(tf): validate flattened multi-device op widths

[njzjz-bot]

Fixes #5623

Summary

• derive per-atom descriptor and neighbor widths through a shared checked helper instead of truncating integer division
• reject partial atom rows, negative local-atom counts, and widths beyond the native kernel integer range
• enforce the se_a four-value descriptor stride and the se_r one-value stride before output allocation
• apply the checks to force, force-gradient, and virial-gradient CPU/GPU dispatch paths
• add focused malformed-input tests for net_deriv, in_deriv, nlist, rij, and descriptor/neighbor stride mismatches

Why existing tests missed this

The existing tests use placeholders whose second dimensions are fixed to the correct nloc * ndescrpt, nloc * ndescrpt * 3, and nloc * nnei widths. Feeding a differently sized array is therefore rejected by TensorFlow before the custom op executes. The tests also cover only valid descriptor layouts, so the release-build paths where assert is compiled out were never exercised.

The new tests construct malformed tensors directly at the op boundary. They verify InvalidArgument is returned by the custom op before output allocation or raw CPU/GPU pointer dispatch.

Validation

• CUDA TensorFlow op build: cmake --build source/build --target deepmd_op
• isolated editable TensorFlow build with current libdeepmd_op and libop_grads
• pytest source/tests/tf/test_multi_device_shape_validation.py -v (6 validation cases plus TensorFlow test_session, all passed)
• existing happy paths: TestProdForce::test_prod_force, TestProdForceGrad::test_prod_force_grad, and TestProdVirialGrad::test_prod_virial_grad (3 passed)
• ruff format .
• ruff check .
• clang-format check on changed C++ files

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

Summary by CodeRabbit

• Bug Fixes

  • Added validation for flattened tensor dimensions in multi-device force and virial operations.
  • Invalid widths, non-divisible dimensions, and inconsistent descriptor or neighbor sizes now return clear errors instead of proceeding to kernel execution.
  • Improved handling of zero-width and oversized per-atom feature dimensions.

• Tests

  • Added regression coverage for malformed shapes across forward and gradient operations.

[coderabbitai]

Warning

Review limit reached

@njzjz, you've reached your PR review limit, so we couldn't start this review.

Next review available in: 32 minutes

Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available.
You're only billed for reviews past your plan's rate limits ($0.25/file).

How can I continue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews.

How do review limits work?

CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability.

For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window.

Please refer docs for additional details.

Review details

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro Plus

Run ID: f1661323-e4b6-467e-b7b5-e8666268552b

📥 Commits

Reviewing files that changed from the base of the PR and between 4cf48fd and e736a90.

📒 Files selected for processing (1)

• source/tests/tf/test_multi_device_shape_validation.py

📝 Walkthrough

Walkthrough

Adds shared flattened-width validation for TensorFlow multi-device custom ops, applies descriptor and neighbor consistency checks to force and virial kernels, and adds regression tests for malformed dimensions.

Changes

Multi-device shape validation

Layer / File(s)	Summary
Per-atom width validation source/op/tf/custom_op.h	Adds GetPerAtomWidth to validate nloc, flattened-width divisibility, zero-width shapes, and int bounds.
Force operator validation source/op/tf/prod_force_multi_device.cc, source/op/tf/prod_force_grad_multi_device.cc	Uses the shared helper and validates descriptor widths against neighbor widths for forward and gradient force operations.
Virial gradient validation source/op/tf/prod_virial_grad_multi_device.cc	Uses shared width validation and enforces descriptor-to-neighbor width relationships.
Regression tests source/tests/tf/test_multi_device_shape_validation.py	Adds tests for invalid divisibility, descriptor counts, descriptor strides, neighbor widths, and virial dimensions.

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

Sequence Diagram(s)

sequenceDiagram
  participant TensorFlowCustomOp
  participant GetPerAtomWidth
  participant RawKernel
  TensorFlowCustomOp->>GetPerAtomWidth: validate flattened width and nloc
  GetPerAtomWidth-->>TensorFlowCustomOp: return width or InvalidArgument
  TensorFlowCustomOp->>TensorFlowCustomOp: validate descriptor/neighbor relationship
  TensorFlowCustomOp->>RawKernel: dispatch validated tensors

Loading

Possibly related PRs

• deepmodeling/deepmd-kit#5600: Introduces the TensorFlow compatibility validation wrapper used by GetPerAtomWidth.

Suggested labels: Core

Suggested reviewers: njzjz

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title is concise and accurately summarizes the main change: validating flattened multi-device op widths.
Linked Issues check	✅ Passed	The changes add pre-kernel width validation across force, force-grad, and virial-grad paths and include regression tests for invalid widths.
Out of Scope Changes check	✅ Passed	The helper, validations, and tests are all directly related to the multi-device shape-validation fix.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ 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 32.35294% with 23 lines in your changes missing coverage. Please review.
✅ Project coverage is 78.62%. Comparing base (1c5ff68) to head (e736a90).
⚠️ Report is 44 commits behind head on master.

Files with missing lines	Patch %	Lines
source/op/tf/custom_op.h	31.25%	7 Missing and 4 partials ⚠️
source/op/tf/prod_force_grad_multi_device.cc	33.33%	4 Missing ⚠️
source/op/tf/prod_force_multi_device.cc	33.33%	4 Missing ⚠️
source/op/tf/prod_virial_grad_multi_device.cc	33.33%	4 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5827      +/-   ##
==========================================
+ Coverage   78.57%   78.62%   +0.04%     
==========================================
  Files        1049     1054       +5     
  Lines      120650   121779    +1129     
  Branches     4348     4409      +61     
==========================================
+ Hits        94806    95746     +940     
- Misses      24285    24462     +177     
- Partials     1559     1571      +12     

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

[njzjz]

Possible reviewers based on changed lines, exact file history, and exact-file review history:

• @wanghan-iapcm — 4 reviews on exact changed files (source/op/tf/custom_op.h).

No review request was made automatically.

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

[wanghan-iapcm]

Nicely done — returning a Status via OP_REQUIRES_OK (rather than throwing) means the error surfaces as a clean InvalidArgumentError and avoids the safe_compute→errors::Internal downgrade. Two notes; the test-coverage one is inline below.

Incomplete fix — the forward virial op is left with the same latent OOB. ProdVirialSeAOp / ProdVirialSeROp in prod_virial_multi_device.cc still use the unchecked truncating dim_size(1) / nloc with no GetPerAtomWidth / nnei*4 == ndescrpt guard:

deepmd-kit/source/op/tf/prod_virial_multi_device.cc

Lines 75 to 79 in 4cf48fd

	int nall = natoms[1];
	int nnei = nloc > 0 ? nlist_tensor.shape().dim_size(1) / nloc : 0;
	int nframes = net_deriv_tensor.shape().dim_size(0);
	int ndescrpt = nloc > 0 ? net_deriv_tensor.shape().dim_size(1) / nloc : 0;
	// check the sizes

. This is reachable: the kernel prod_virial_a_cpu recomputes ndescrpt = 4 * nnei internally and indexes with that, so the op's existing nloc*ndescrpt*3 == in_deriv_width / nloc*nnei*3 == rij_width checks (which only tie the buffers to the truncated widths, never ndescrpt against nnei) don't prevent a malformed net_deriv width from driving an out-of-bounds read. It runs on every inference/training step. Issue #5623 enumerated only the three ops you fixed, but the forward virial shares the identical pattern — worth applying the same helper + stride guard here (and a matching test).

[wanghan-iapcm]

All four gaps are closed, and I was able to verify it by running rather than reading, which is worth spelling out because it makes the result unambiguous.

The TensorFlow custom ops installed in my environment predate this change -- neither libdeepmd_op.so nor libop_grads.so contains the "should be divisible by nloc" string -- so running the new suite against them is a genuine pre-fix run. Eight of the twelve collected cases fail there, including every one of the four I asked for:

gap I raised	new test	unpatched
GetPerAtomWidth nloc < 0	test_negative_nloc_is_rejected	fails
nloc == 0	test_zero_nloc_rejects_nonempty_flattened_width	fails
prod_force_se_r_grad untested	test_prod_force_r_grad_rejects_partial_nlist_atom	fails
prod_virial_se_r_grad untested	test_prod_virial_r_grad_rejects_partial_net_deriv_atom	fails

Including test_zero_nloc_accepts_empty_flattened_widths was a good call -- the zero-nloc guard has a legitimate accepting case, and without it a future change could satisfy every rejection test by simply refusing all zero-nloc input.

Two of the remaining cases, test_prod_force_rejects_in_deriv_width_mismatch and test_prod_virial_grad_rejects_rij_width_mismatch, already pass on the unpatched ops, because an existing check rejects those inputs. That is not a problem and I am not asking for anything -- they are regression guards for behaviour that already worked rather than demonstrations of the bug -- but it is worth knowing which of these tests would actually catch a revert of this PR and which would not.

Agreed on leaving the width > INT_MAX branch alone; it is not practically constructible through the op API and a test for it would cost more than it protects.

Approving.