fix(tf): validate GPU tabulation sizes

[njzjz-bot]

Summary

• validate 1 <= last_layer_size <= 1024 before every TensorFlow CUDA/ROCm tabulation wrapper call
• move the four existing grad-gradient checks ahead of their kernel launches
• cover all SeA, SeAtten, SeT, and SeR forward/gradient/grad-gradient GPU paths, including zero width

Why existing tests missed this

There were no direct TensorFlow tabulation-op GPU tests. Existing PyTorch tabulation coverage uses normal small embedding widths and does not execute these TensorFlow kernels, while real compressed models are also normally well below the device limit. The four existing grad-gradient checks looked like validation but ran only after the GPU wrapper had already launched, and forward/first-gradient paths had no checks at all.

Validation

• cmake --build source/build --target deepmd_op -j2
• constructed all 12 direct-op regression graphs against the fresh build-tree TensorFlow extension
• pytest source/tests/tf/test_tabulate_gpu_size_validation.py source/tests/tf/test_tabulate.py -q (2 passed, 2 GPU-only tests skipped locally because no usable GPU is exposed)
• ruff format .
• ruff check .
• clang-format dry run with --Werror
• git diff --check

Fixes #5654

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

Summary by CodeRabbit

• Bug Fixes

  • GPU tabulation operations now reject invalid layer sizes below 1 or above 1024.
  • Invalid configurations return a consistent error message before GPU execution begins.

• Tests

  • Added GPU coverage for oversized and zero-width layer size validation across supported tabulation 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: 2626767f-29c4-408b-9109-fc697d38f892

📥 Commits

Reviewing files that changed from the base of the PR and between 6c3b985 and 6ce049d.

📒 Files selected for processing (2)

• source/op/tf/tabulate_multi_device.cc
• source/tests/tf/test_tabulate_gpu_size_validation.py

---

📝 Walkthrough

Walkthrough

The TensorFlow GPU tabulation wrappers now validate last_layer_size in the range 1–1024 before launching CUDA/ROCm kernels. GPU-gated tests cover oversized and zero values across forward, gradient, and grad-gradient paths.

Changes

GPU tabulation validation

Layer / File(s)	Summary
Shared GPU size validator source/op/tf/tabulate_multi_device.cc	Adds a common validator that returns InvalidArgument unless last_layer_size is between 1 and 1024.
Wrapper prelaunch checks source/op/tf/tabulate_multi_device.cc	All forward, gradient, and grad-gradient GPU wrappers validate the size before launching tabulation kernels.
Invalid-size test coverage source/tests/tf/test_tabulate_gpu_size_validation.py	Adds GPU-gated tests for 1025 across all wrapper paths and 0 for selected paths, requiring the standardized error message.

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

Suggested labels: OP, CUDA, ROCM

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.
Linked Issues check	✅ Passed	GPU validation is added before tabulation kernel launches, rejects <=0 and >1024, and tests cover oversized and zero-width cases.
Out of Scope Changes check	✅ Passed	The changes are limited to the GPU validation fix and matching regression tests.
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 0% with 15 lines in your changes missing coverage. Please review.
✅ Project coverage is 78.31%. Comparing base (6c3b985) to head (6ce049d).

Files with missing lines	Patch %	Lines
source/op/tf/tabulate_multi_device.cc	0.00%	15 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5846      +/-   ##
==========================================
- Coverage   78.58%   78.31%   -0.27%     
==========================================
  Files        1050     1050              
  Lines      120637   120648      +11     
  Branches     4356     4357       +1     
==========================================
- Hits        94801    94488     -313     
- Misses      24278    24599     +321     
- Partials     1558     1561       +3     

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