fix(tabulate): handle empty SE-A neighbors

[njzjz-bot]

Root cause

The shared SE-A CPU and GPU implementations assumed that every local atom had a positive neighbor capacity. They read the last element with an index equivalent to ii * nnei + nnei - 1 before entering the neighbor loop. The framework wrappers guard nloc <= 0, but not nnei <= 0, so a valid empty neighbor axis reaches the native implementation and turns that index into -1.

Forward also needs more than an early return because the descriptor does not contain the empty neighbor dimension and may be allocated uninitialized. Its empty reduction is mathematically zero.

Introduction and history

The non-empty-neighbor assumption dates to 76a1a2a1d from 2021, optimize interface of custom ops, which introduced the shared tabulation loop and the last-neighbor lookup. Later CUDA optimization and the CUDA/ROCm merge preserved the same contract. This is therefore a long-standing missing edge case, not a recent regression.

Test gap

The existing SE-A C++ tests always used a positive nnei. They covered normal forward and derivative values but never exercised nloc > 0, nnei == 0, so the pre-loop dereference and the need to initialize descriptor-shaped outputs were not visible.

Fix

• zero the CPU and GPU forward descriptor and return when nnei <= 0;
• return from first-order gradients before touching empty tensor pointers;
• zero descriptor-shaped second-order outputs before returning;
• add CPU and GPU regression coverage for forward, grad, and grad-grad.

Validation

• CUDA 12.4 build of runUnitTests_lib;
• TestTabulateSeA.empty_neighbors_cpu;
• TestTabulateSeA.empty_neighbors_gpu on an NVIDIA GeForce RTX 5090 allocation;
• ruff format .;
• ruff check ..

Fixes #5892

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

Summary by CodeRabbit

• Bug Fixes

  • Improved CPU and GPU tabulation when no neighbors are present.
  • Empty-neighbor reductions now return safely with zeroed outputs.
  • First- and second-order gradient calculations avoid invalid processing and produce correct results for empty inputs.
  • Prevented unnecessary processing of missing neighbor data, improving robustness for edge-case inputs.

• Tests

  • Added CPU and GPU coverage for empty-neighbor handling, including gradient and second-gradient calculations.

[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: 643dd67d-5e47-4a68-a54f-d6730070314b

📥 Commits

Reviewing files that changed from the base of the PR and between 21ef773 and be42a5f.

📒 Files selected for processing (3)

• source/lib/src/gpu/tabulate.cu
• source/lib/src/tabulate.cc
• source/lib/tests/test_tabulate_se_a.cc

🚧 Files skipped from review as they are similar to previous changes (3)

• source/lib/src/tabulate.cc
• source/lib/tests/test_tabulate_se_a.cc
• source/lib/src/gpu/tabulate.cu

---

📝 Walkthrough

Walkthrough

SE-A CPU and GPU tabulation entrypoints now handle nnei <= 0 without neighbor indexing or kernel launches. Forward and second-order outputs are zero-initialized as needed, and CPU/GPU tests cover empty-neighbor inputs with null neighbor-dependent pointers.

Changes

SE-A empty-neighbor handling

Layer / File(s)	Summary
CPU empty-neighbor guards source/lib/src/tabulate.cc	CPU forward, gradient, and second-order gradient entrypoints return early when nnei <= 0, preserving zero-initialized outputs.
GPU empty-neighbor guards source/lib/src/gpu/tabulate.cu	GPU forward and second-order paths zero outputs and return for empty reductions. The first-order path skips its kernel launch.
Empty-neighbor validation source/lib/tests/test_tabulate_se_a.cc	CPU and GPU tests verify zero outputs and safe handling of null neighbor-dependent pointers for forward, gradient, and second-order gradient calls.

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

Suggested reviewers: wanghan-iapcm

🚥 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 and concisely describes the fix for empty SE-A neighbors.
Linked Issues check	✅ Passed	The changes address issue #5892 by handling zero neighbors, initializing outputs, and adding CPU and GPU regression tests.
Out of Scope Changes check	✅ Passed	All changes support the linked issue by fixing empty-neighbor handling and adding focused 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 85.71429% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.16%. Comparing base (cf3e6f1) to head (be42a5f).
⚠️ Report is 4 commits behind head on master.

Files with missing lines	Patch %	Lines
source/lib/src/tabulate.cc	50.00%	0 Missing and 3 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5934      +/-   ##
==========================================
- Coverage   79.41%   79.16%   -0.25%     
==========================================
  Files        1071     1071              
  Lines      124844   124865      +21     
  Branches     4531     4536       +5     
==========================================
- Hits        99144    98850     -294     
- Misses      24079    24390     +311     
- Partials     1621     1625       +4     

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

The fix is correct and correctly scoped, and the part that took judgement rather than a reflex is right: an early return alone would have been wrong for the forward and second-order paths. out and dz_dy are sized nloc * 4 * last_layer_size independently of nnei, and PyTorch allocates them via torch::empty/empty_like while TF uses allocate_output, so returning without materializing the zero reduction would have handed back uninitialized memory rather than a correct empty result. The first-order outputs genuinely are neighbor-shaped — dy_dem_x is nloc*nnei, dy_dem is nloc*nnei*4, dy_dtwo is nloc*nnei*last_layer_size — so returning ahead of their memsets is safe. I checked dy_dtwo specifically, since a descriptor-shaped output hiding in the gradient path would have made that early return a bug; it mirrors two_embed's neighbor shape at the caller.

Scope is right too. The sentinel read em_x[ii * nnei + nnei - 1] exists only in the SE-A family (tabulate.cc:178,266,367 on master, and the three GPU counterparts). SE-T and SE-R either memset their outputs before looping or write unconditionally, and se_t_tebd already guards nnei_i/nnei_j from #4992, so leaving them alone is correct rather than an omission.

On the test: the existing SE-A tests all use a positive nnei, so the untested cell was nloc > 0 x nnei == 0 — the empty case had only ever been probed on the atom axis, never the neighbor axis. The new test crosses it and genuinely fails pre-fix: compiled against master's tabulate.cc it segfaults on em_x[-1], and passes against this branch. The forward and grad-grad halves also pre-fill with 1.0 and assert exact zeros, so they pin the zeroing behaviour and not merely absence of a crash. That is a real regression test.

Two corrections to the description, neither affecting the code. The origin is not 76a1a2a1d: git show 76a1a2a1d^:source/lib/include/CustomeOperation.h already contains FPTYPE ago = in[ii * nnei + nnei - 1];, so that commit relocated the lookup rather than introducing it. It came from b5166a818 ("model compression", 2021-02-03) three weeks earlier. And "the framework wrappers guard nloc <= 0" is not the case — grep -n "nloc <= 0" source/op/tf/tabulate_multi_device.cc source/op/pt/tabulate_multi_device.cc returns nothing. The nloc guards live in the GPU entry points of the shared lib, added by 8ce9884a3 (#808) for the CUDA launch-configuration constraint that <<<0, ...>>> is an error, not as an input contract; the CPU entry points still have no nloc guard at all. The placement you chose is still the right one — #4992 set that precedent when it added lib-level nnei_i/nnei_j guards for se_e3_tebd — so this only affects the rationale, not the outcome.

One minor thing worth a follow-up rather than a round trip here. In tabulate_fusion_se_a_gpu and tabulate_fusion_se_a_grad_grad_gpu the new nnei <= 0 branch issues gpuMemset and gpuDeviceSynchronize before the DPErrcheck(gpuGetLastError()) that every entry point in this file otherwise runs first. A sticky error left by an earlier kernel would then be reported against this memset and blamed on the empty-neighbor path. It only affects which line gets named in an already-failing program, but moving the guard below the two existing DPErrcheck lines would cost nothing.

Finally, on severity: reaching nnei == 0 from Python needs sum(sel) == 0, i.e. every type excluded, and #5892 notes a bare CUDA probe did not surface a driver error. The out-of-bounds read is real at the library API level regardless — worth fixing as hardening, and the guard is free on the hot path — but I would not describe it as a live crash users are hitting.

[njzjz]

Resolved the master conflict and addressed the review mirror note in be42a5f.

• Merged current origin/master and retained both required test headers in test_tabulate_se_a.cc.
• Moved the existing gpuGetLastError/synchronization checks ahead of the new empty-neighbor memset branches in the SE-A forward and grad-grad GPU entrypoints, so a sticky prior error is reported at the established check instead of being misattributed to the empty-neighbor path.

Validation: ruff format ., ruff check ., clang-format --dry-run --Werror, and diff checks passed; the resolved C++ test file passed a standalone syntax check; the backend-neutral deepmd C++ library rebuilt successfully with the merged CPU tabulation source. CUDA/ROCm execution remains CI-only in this environment.

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