[Code scan] SE-Attention ops can read and write past mismatched tensor shapes

[njzjz-bot]

This issue comes from a focused memory-safety review at commit 21b73ab303581f467bba0cae101a880c35cc2d82.

Problem

The TensorFlow and PyTorch SE-Attention wrappers validate tensor rank but do not validate the dimensions or element counts that the native implementation indexes.

For PyTorch:

• forward checks only that em_x and two_embed are rank 2 and em is rank 3
• nloc and nnei are taken exclusively from em
• backward allocates dy_dtwo with zeros_like(two_embed)

The native code nevertheless indexes or clears exactly nloc * nnei * last_layer_size two_embed/ dy_dtwo elements:

• CPU reads two_embed[ii * nnei * last_layer_size + jj * last_layer_size + kk]
• CPU backward clears that full computed size
• GPU backward writes the same computed range

TensorFlow has the same mismatch: it checks only ranks, while the backward output is allocated using the unvalidated two_embed shape.

The same class of problem applies to em_x, the final dimension of em (assumed to be 4), dy, descriptor, table_info, and the table size.

Reproduction

The probes were run inside an srun --partition=main --gres=gpu:5090:1 allocation.

A PyTorch tensor with logical shape [1, 1] was passed as two_embed for em.shape == [1, 3, 4]. Its backing storage contained three values [0, 10, 20]. Although only the first value belongs to the tensor view, forward consumed all three:

logical_two_shape (1, 1) logical_numel 1
backing_storage [0.0, 10.0, 20.0]
output [6.6000000000000005, 33.0, 0.0, 0.0]

A C++ canary probe then emulated the one-element backward output followed by two adjacent canary values. The canaries started as 12345 and 67890; after tabulate_fusion_se_a_grad_cpu they had both been overwritten:

logical_output=1.2 canary0=1.2 canary1=1.2

Thus the missing validation causes both reads beyond a tensor's logical boundary and writes beyond the allocated output shape. On GPU the corresponding write can become an illegal-address error; on CPU it can silently corrupt another allocation or eventually segfault.

Impact

Normal descriptor code currently constructs matching shapes, so this is primarily reachable through direct/raw custom-op calls or an upstream shape bug. However, custom ops must reject malformed inputs rather than corrupt process memory.

Suggested fix

Before taking raw pointers or launching the native kernels, validate at least:

• em.shape == [nloc, nnei, 4];
• em_x.shape == [nloc, nnei];
• two_embed.numel() == nloc * nnei * last_layer_size;
• forward/backward descriptor and cotangent shapes;
• table_info.numel() >= 5;
• the table contains every interval and feature coefficient that table_info and last_layer_size can address;
• required device placement is consistent.

Apply equivalent checks to the related SE-A, SE-T, and SE-R raw ops.

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