ScatterND in a dynamic-shape engine silently drops updates beyond ~N when the update count M exceeds the output size N (static engine correct)

[braindevices]

ScatterND in a dynamic-shape engine silently drops updates beyond ~N when the update count M exceeds the
output size N (static engine correct)

Environment: TensorRT 11.2.1.2 (tensorrt-cu13 Python wheels), CUDA runtime 13.0.96, driver 580.142; Quadro RTX
4000 (sm_75) and RTX 4000 SFF Ada Generation (sm_89); Linux (AlmaLinux 9.8), Python 3.14; ONNX opset 18, strongly
typed network, one optimization profile.

Description: A single ONNX ScatterND (data (N,), indices (M, 1), updates (M,)) built into an engine with a
dynamic-shape optimization profile applies only about the first N entries of the update list when M > N; later
entries are silently not written. The same graph built as a static engine, onnxruntime, and the ONNX reference
evaluator apply all updates. With M <= N the dynamic engine is correct, even with the same duplicate indices.
Making either N or M dynamic is enough. fp32 and fp16 behave the same; both GPUs above are affected.

The update list in the repro contains duplicate indices (all non-test updates write index N - 1), which the ONNX
spec says ScatterND indices should not have, because the iteration order is unspecified. The updates that get
dropped, however, target indices that occur only once, so any iteration order would write them; the result differs
between static and dynamic engines of the same graph, and depends on M vs N. There is no build- or run-time
warning. If duplicates make this undefined in TensorRT, a warning or an error would help; we lost masks in
production-like output without any signal.

Steps to reproduce: run the script below (needs tensorrt, onnx, numpy, and torch with CUDA for device
buffers).

"""ScatterND in a dynamic-shape engine drops
updates past about N when M > N (with M <= N, the same duplicates are fine).

One ScatterND: data (N,) zeros, indices (M, 1), updates (M,). 1000 "real" updates with unique indices 0..999
(values 1..1000) are spread evenly over the update list; every other update writes index N-1. Expected (ONNX
reference / onnxruntime / a static TensorRT engine): out[0:1000] == 1..1000.
Needs: tensorrt, onnx, numpy, torch (CUDA, only for device buffers).
"""
import numpy as np
import tensorrt as trt
import torch

from onnx import TensorProto
from onnx import helper
from onnx.reference import ReferenceEvaluator

N, M, REAL = 91585, 144000, 1000


def model(dynamic: bool) -> bytes:
    n, m = ("N", "M") if dynamic else (N, M)
    g = helper.make_graph(
        [helper.make_node("ScatterND", ["data", "indices", "updates"], ["out"])], "scatternd",
        [helper.make_tensor_value_info("data", TensorProto.FLOAT, [n]),
         helper.make_tensor_value_info("indices", TensorProto.INT64, [m, 1]),
         helper.make_tensor_value_info("updates", TensorProto.FLOAT, [m])],
        [helper.make_tensor_value_info("out", TensorProto.FLOAT, [n])])
    mdl = helper.make_model(g, opset_imports=[helper.make_opsetid("", 18)])
    mdl.ir_version = 9
    return mdl.SerializeToString()


def run(dynamic: bool, data, indices, updates) -> np.ndarray:
    logger = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(logger)
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))
    parser = trt.OnnxParser(network, logger)
    assert parser.parse(model(dynamic))
    config = builder.create_builder_config()
    if dynamic:
        p = builder.create_optimization_profile()
        p.set_shape("data", (N,), (2 * N,), (6 * N,))
        p.set_shape("indices", (M, 1), (2 * M, 1), (15 * M, 1))
        p.set_shape("updates", (M,), (2 * M,), (15 * M,))
        config.add_optimization_profile(p)
    engine = trt.Runtime(logger).deserialize_cuda_engine(builder.build_serialized_network(network, config))
    ctx = engine.create_execution_context()
    bufs = {"data": torch.from_numpy(data).cuda(), "indices": torch.from_numpy(indices).cuda(),
            "updates": torch.from_numpy(updates).cuda(), "out": torch.empty(len(data), device="cuda")}
    for name in ("data", "indices", "updates"):
        ctx.set_input_shape(name, tuple(bufs[name].shape))
    for name, t in bufs.items():
        ctx.set_tensor_address(name, t.data_ptr())
    stream = torch.cuda.Stream()
    ctx.execute_async_v3(stream.cuda_stream)
    stream.synchronize()
    return bufs["out"].cpu().numpy()


def inputs(n: int) -> tuple[np.ndarray, ...]:
    """1000 real updates (unique indices 0..999) spread over the M-long update list; the rest write index n - 1."""
    positions = np.linspace(0, M - 1, REAL).astype(np.int64)
    indices = np.full((M, 1), n - 1, np.int64)
    indices[positions, 0] = np.arange(REAL)
    updates = np.full(M, -1.0, np.float32)
    updates[positions] = np.arange(1, REAL + 1)
    return np.zeros(n, np.float32), indices, updates, positions


expected = np.arange(1, REAL + 1, dtype=np.float32)
print(f"TensorRT {trt.__version__}, {torch.cuda.get_device_name()}; M = {M} updates")
data, indices, updates, positions = inputs(N)
(reference,) = ReferenceEvaluator(model(dynamic=False)).run(None, {"data": data, "indices": indices, "updates": updates})
print(f"  N = {N} (M > N), ONNX reference: {int((reference[:REAL] == expected).sum())}/{REAL} real updates applied")
for dynamic, n in ((False, N), (True, N), (True, 2 * N)):  # the last: the same duplicates, but M <= N
    data, indices, updates, positions = inputs(n)
    ok = run(dynamic, data, indices, updates)[:REAL] == expected
    first_dropped = positions[~ok].min() if (~ok).any() else None
    print(f"  N = {n} ({'M > N' if n < M else 'M <= N'}), {'dynamic' if dynamic else 'static '} engine: {int(ok.sum())}/{REAL} real updates applied"
          + ("" if first_dropped is None else f"; first dropped at update-list position {first_dropped}"))

Expected: every line reports 1000 / 1000 real updates applied.

Actual (both GPUs; the cutoff differs slightly):

TensorRT 11.2.1.2, NVIDIA RTX 4000 SFF Ada Generation; M = 144000 updates
  N = 91585 (M > N), ONNX reference: 1000/1000 real updates applied
  N = 91585 (M > N), static  engine: 1000/1000 real updates applied
  N = 91585 (M > N), dynamic engine: 636/1000 real updates applied; first dropped at update-list position 91675
  N = 183170 (M <= N), dynamic engine: 1000/1000 real updates applied
TensorRT 11.2.1.2, Quadro RTX 4000; M = 144000 updates
  N = 91585 (M > N), ONNX reference: 1000/1000 real updates applied
  N = 91585 (M > N), static  engine: 1000/1000 real updates applied
  N = 91585 (M > N), dynamic engine: 640/1000 real updates applied; first dropped at update-list position 92251
  N = 183170 (M <= N), dynamic engine: 1000/1000 real updates applied

Additional observations: the applied updates are always a prefix of the update list, ending within 0.1-0.7 % of
N, across N = 91,585 - 547,244 and M = 144,000 - 2,112,000 in one profile; only M dynamic (N static) or only N
dynamic (M static) reproduce it too.