[Code scan] SE-A tabulation dereferences `em_x[-1]` when `nnei == 0`

[njzjz-bot]

This issue comes from a focused review of tabulation at commit 21b73ab303581f467bba0cae101a880c35cc2d82.

Problem

The SE-A GPU kernels assume at least one neighbor and read the last em_x element before processing the neighbor loop:

• forward reads em_x[block_idx * nnei + nnei - 1]
• gradient reads the same location while finding the shared breakpoint
• grad-grad reads it unconditionally

When nloc > 0 and nnei == 0, each address is em_x[-1]. The wrappers guard only nloc <= 0, so they still launch these kernels for empty-neighbor tensors.

The CPU implementations mirror the same assumption by evaluating em_x[ii * nnei + nnei - 1] before their loops.

Impact

Shapes such as em_x=[nloc, 0] and em=[nloc, 0, 4] are representable by the TensorFlow and PyTorch wrappers. They can trigger an invalid pointer read, a CUDA/HIP illegal-address error, or undefined output. A simple early return is not sufficient for forward because some wrappers allocate the descriptor with empty; the mathematically correct empty reduction must be initialized to zero.

Validation note

The source-level out-of-bounds access is deterministic from the index expression. A bare CUDA probe on an NVIDIA GeForce RTX 5090 did not surface a driver error, and CUDA 12.4 compute-sanitizer on this system reports Error: Device not supported. for that GPU, so there is no sanitizer trace attached.

Suggested fix

Handle nnei <= 0 explicitly in CPU and GPU entry points:

• zero the forward descriptor and return;
• leave/initialize first- and second-order gradient outputs to zero and return;
• add CPU and GPU tests for nloc > 0, nnei == 0.

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

[njzjz-bot]

A reachability/scope note after checking the higher-level neighbor-list paths:

• In normal training and inference, nnei is the fixed neighbor capacity, typically sum(sel), rather than the number of neighbors actually found inside rcut.
• Therefore, an isolated atom with no physical neighbors does not trigger this issue when sum(sel) > 0; its neighbor row is still padded to that width.
• A per-type value such as sel[i] == 0 is also harmless as long as the total sum(sel) remains positive.
• The failing case requires nloc > 0 and a zero total capacity, for example all sel entries being zero, sel: auto producing all zeros for data with no neighbors inside the cutoff, or a direct custom-op call with shapes such as [nloc, 0].

I did not find a high-level validation rule that rejects a zero total sel; additionally, the automatic selection logic preserves zero because rounding zero to a multiple of four still gives zero. Thus this appears to be a low-frequency but reachable robustness bug, rather than a common production failure.

The exact em_x[nnei - 1] pattern is confined to the SE-A tabulation implementation, but the implementation is shared by the TensorFlow legacy/SE-A/SE-Atten registrations and the PyTorch SE-A/SE-Atten autograd paths. Handling the empty reduction in the shared CPU/GPU entry points would cover those aliases. Alternatively, the public boundary would need to reject sum(sel) == 0 explicitly; leaving the case as an out-of-bounds access is unsafe.

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