fix(ops): validate tabulation tensor shapes

[njzjz-bot]

Root cause

The TensorFlow and PyTorch tabulation wrappers checked mostly tensor rank, then derived nloc, neighbor dimensions, and output sizes from other tensors. The native CPU and GPU code indexes flat buffers using those derived sizes. A rank-correct but shorter em_x, em, two_embed, cotangent, descriptor, table_info, or table could therefore be read or written beyond its logical allocation.

Introduction and history

This is an accumulated public-boundary validation gap rather than one recent native-kernel regression:

• the PyTorch rank-only wrappers were added by 54aff9571, PR pt: add 4 tabulate_fusion op #3877, in 2024;
• the TensorFlow SE-Attention wrapper was added by dddba6e8c, PR Enable model compression for se_atten #2532, in 2023;
• the underlying TensorFlow tabulation wrappers and raw-buffer contract predate them, including 76a1a2a1d and 472d7c4e5.

Each wrapper trusted shapes produced by the normal descriptor path, but raw custom-op calls and upstream shape bugs can violate that trust boundary.

Test gap

Existing PyTorch tests covered matched forward, backward, and second-order inputs only. TensorFlow had model-compression tests for valid graphs but no focused raw-op rejection tests. They also did not explicitly preserve the legal flattened rank-2 em_x layout used by TensorFlow compression, which was important while tightening validation.

Fix

• add shared validation for finite table metadata, positive strides, generated spline-row capacity, and overflow-safe element counts;
• validate dtype, device, contiguity, rank, and required element counts in the public PyTorch SE-A, SE-Attention, SE-T, SE-T-TEBD, and SE-R forward ops;
• validate forward, grad, and grad-grad input and cotangent shapes for the TensorFlow SE-A, SE-Attention, SE-T, and SE-R kernels;
• preserve legal flattened rank-2 layouts by validating logical element counts rather than requiring one exact rank-2 shape;
• reject malformed buffers before taking raw pointers or launching native kernels.

Validation

• CPU editable build with TensorFlow 2.21 and PyTorch custom ops;
• 42 focused PyTorch and TensorFlow tabulation tests passed, plus 6 subtests;
• valid TensorFlow compressed-model test_1frame scenarios passed for SE-A, SE-Attention, SE-T, and SE-R;
• explicit rejection coverage for mismatched em_x, truncated em, short two_embed, short table_info, short table storage, and mismatched gradient shapes;
• ruff format .;
• ruff check ..

Fixes #5895

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

Summary by CodeRabbit

• Bug Fixes

  • Added stronger validation for tabulation inputs, including tensor shapes, data types, device placement, memory layout, and table sizes.
  • Prevented invalid metadata, overflow conditions, and insufficient coefficient data from reaching computation.
  • Improved error reporting for malformed TensorFlow and PyTorch tabulation inputs.
  • Confirmed support for compatible flattened descriptor layouts.

• Tests

  • Added coverage for invalid shapes, short buffers, malformed metadata, gradient dimensions, and accepted flattened inputs across TensorFlow and PyTorch.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro Plus

Run ID: c4cf59b3-f158-4f83-ae57-8c06eeda201d

📥 Commits

Reviewing files that changed from the base of the PR and between 4f827cc and 788f1f5.

📒 Files selected for processing (5)

• source/lib/include/tabulate_validation.h
• source/op/pt/tabulate_multi_device.cc
• source/op/tf/tabulate_multi_device.cc
• source/tests/pt/test_tabulate_fusion_se_atten.py
• source/tests/tf/test_tabulate_shape_validation.py

---

📝 Walkthrough

Walkthrough

Shared checked arithmetic validates tabulation metadata and required buffer sizes. PyTorch and TensorFlow tabulation operations now validate tensor contracts before native dispatch, with new tests covering malformed shapes, flattened layouts, and undersized buffers.

Changes

Tabulation validation

Layer / File(s)	Summary
Shared table-size validation source/lib/include/tabulate_validation.h	Adds checked multiplication, metadata validation, required row calculation, and coefficient element-count validation.
PyTorch pre-dispatch checks source/op/pt/tabulate_multi_device.cc, source/tests/pt/test_tabulate_fusion_se_atten.py	Validates devices, dtypes, contiguity, shapes, and table capacity before SeA, SeT, SeR, and SeT-TEBD dispatch; adds malformed-input coverage.
TensorFlow kernel validation source/op/tf/tabulate_multi_device.cc, source/tests/tf/test_tabulate_shape_validation.py	Centralizes table, tensor, descriptor, gradient, and derived-size checks across forward and gradient kernels; adds shape and buffer validation tests.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Caller
  participant TabulationWrapper
  participant ValidationHelpers
  participant NativeKernel
  Caller->>TabulationWrapper: provide table and operation tensors
  TabulationWrapper->>ValidationHelpers: validate shapes and required capacity
  ValidationHelpers-->>TabulationWrapper: accept or report validation error
  TabulationWrapper->>NativeKernel: dispatch validated inputs

Loading

Possibly related PRs

• deepmodeling/deepmd-kit#5829: Adds overlapping PyTorch tabulation device-contract validation.
• deepmodeling/deepmd-kit#5846: Overlaps with TensorFlow tabulation validation for derived size constraints.

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 clearly matches the main change: adding validation for tabulation tensor shapes.
Linked Issues check	✅ Passed	The PR implements the required shape, size, device, and storage validations for the affected tabulation ops and adds rejection tests.
Out of Scope Changes check	✅ Passed	The diff stays focused on tabulation validation helpers, op checks, and matching tests with no obvious unrelated changes.
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 47.84483% with 121 lines in your changes missing coverage. Please review.
✅ Project coverage is 78.90%. Comparing base (4f827cc) to head (788f1f5).

Files with missing lines	Patch %	Lines
source/op/tf/tabulate_multi_device.cc	41.04%	64 Missing and 15 partials ⚠️
source/lib/include/tabulate_validation.h	51.02%	19 Missing and 5 partials ⚠️
source/op/pt/tabulate_multi_device.cc	63.26%	18 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5932      +/-   ##
==========================================
- Coverage   79.21%   78.90%   -0.31%     
==========================================
  Files        1069     1070       +1     
  Lines      124070   124302     +232     
  Branches     4522     4551      +29     
==========================================
- Hits        98278    98079     -199     
- Misses      24171    24584     +413     
- Partials     1621     1639      +18     

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