🐛 [Bug] index_put with a None index materialises one index per element, failing on a runtime-known extent

[SrivastavaKshitij]

Bug Description

A None entry in an index_put index list leaves that dimension unindexed, and Torch
broadcasts values across it. ScatterND has no equivalent of an unindexed dimension, so
index_put_converter compensates by materialising one index row per element of the free
dimensions -- N * F_volume rows, where N is the index length -- and separately reshaping
and expanding values to (N,) + F_shape_values. That construction is only correct when
those extents are known at build time. When the index length is only known at runtime, the
per-dimension broadcast check accepts a mismatch instead of rejecting it, and the emitted
expand ends up reading past the end of an axis.

py/torch_tensorrt/dynamo/conversion/impl/select.py :: index_put_converter

    expected_shape = (N,) + tuple(F_shape_values)                       # line 970
    ...
        else:
            # Discontinuous case (K > 1 or K == 0)
            values_shape_padded = [1] * (
                len(expected_shape) - len(values.shape)
            ) + list(values.shape)
            broadcast_shape = []
            for exp_dim, val_dim in zip(expected_shape, values_shape_padded):
                if val_dim == DYNAMIC_DIM or exp_dim == DYNAMIC_DIM:     # line 1065
                    broadcast_shape.append(-1)
                elif val_dim == 1 or exp_dim == val_dim:
                    broadcast_shape.append(exp_dim)
                else:
                    raise ValueError(
                        f"Cannot broadcast {values.shape} to {expected_shape}"
                    )

            values_expanded = impl.slice.expand(
                ctx,
                target,
                source_ir,
                f"{name}_expand_values",
                values,
                expected_shape,
            )

For data[:, idx] = values on rank-2 data with 1-D values, K == 1 and len(F) == 1, so
neither of the earlier branches applies and this else branch (labelled "Discontinuous case")
is the one taken. expected_shape is (N, 32) with N dynamic; values_shape_padded is
[1, -1]. The second pair, (32, -1), is a real mismatch -- the static free-dim extent
against the runtime index length -- but line 1065 fires first and accepts it by appending
-1.

impl.slice.expand is then called with expected_shape. It prepends ones to reach the target
rank, so the rank-1 values of runtime extent N becomes (1, N), and its stride loop takes
the branch commented # No broadcasting is happening. The output should have the same size as input at this dimension. for the axis whose input extent is dynamic and whose target is the
static 32. The result is an ISliceLayer reading 32 elements along an axis whose runtime
extent is only N.

Worth noting: broadcast_shape is computed in that branch but never used -- the expand call
passes expected_shape -- so the loop's only effect is whether it raises, and line 1065 is
what stops it raising.

This fails while BUILDING the engine, not at execution. TensorRT's shape analyzer catches
the contradiction during buildEngineWithConfig:

IBuilder::buildEngineWithConfig: Error Code 4: Internal Error (kOPT values for profile 0
violate shape constraints: [SLICE]-[aten_ops.index_put.default]-[index_put_expand_values]:
ISliceLayer has out of bounds access on axis 1 Condition '<' violated: 31 >= 4.

31 is the last index of the free dimension of extent 32; 4 is the index length at kOPT.

To Reproduce

docker run --rm --gpus all --ipc=host -v "$PWD":/w -w /w \
        nvcr.io/nvidia/pytorch:26.07-py3 python repro.py

repro.py

repro.py

import sys
import traceback

import torch
import torch_tensorrt  # noqa: F401  (registers the "tensorrt" torch.compile backend)


class OneDimValues(torch.nn.Module):
    """data[:, idx] = values with 1-D values, broadcast across the free dim 0."""

    def forward(self, data: torch.Tensor, idx: torch.Tensor) -> torch.Tensor:
        values = idx.to(torch.float32)
        return torch.ops.aten.index_put.default(data, [None, idx], values)


class FullRankValues(torch.nn.Module):
    """data[:, idx] = values with values in torch's layout for data[:, idx]."""

    def forward(self, data: torch.Tensor, idx: torch.Tensor) -> torch.Tensor:
        values = idx.to(torch.float32).reshape(1, -1, 1).expand(32, -1, 128)
        return torch.ops.aten.index_put.default(data, [None, idx], values)


def run_case(
    label: str, model: torch.nn.Module, data: torch.Tensor, idx: torch.Tensor, dynamic: bool
) -> bool:
    """Compiles and runs `model`, marking the index length dynamic if asked."""
    torch._dynamo.reset()
    idx = idx.clone()
    if dynamic:
        torch._dynamo.mark_dynamic(idx, 0)
    model = model.eval().cuda()
    optimized = torch.compile(
        model,
        backend="tensorrt",
        options={"pass_through_build_failures": True, "min_block_size": 1},
    )
    print(f"\n===== {label} =====", flush=True)
    try:
        out = optimized(data, idx)
        torch.testing.assert_close(out, model(data, idx))
    except Exception:  # pylint: disable=broad-except
        traceback.print_exc()
        print(f"----- {label}: FAILED", flush=True)
        return False
    print(f"----- {label}: OK {tuple(out.shape)}, matches eager", flush=True)
    return True


def main(argv: list[str] | tuple[str, ...] = ()) -> int:
    """Runs the two runtime-extent formulations and the static-extent control."""
    del argv
    print(f"torch          {torch.__version__}")
    print(f"torch_tensorrt {torch_tensorrt.__version__}")
    # Severity 4 == VERBOSE. Not needed on the versions tested here - the build-time
    # error is reported at ERROR severity either way - but kept because when this
    # failure surfaces at execution instead, the engine can throw before writing its
    # outputs, in which case the first read reports `_Map_base::at` and python surfaces
    # an unrelated "a Tensor with 0 elements cannot be converted to Scalar".
    torch.ops.tensorrt.set_logging_level(4)

    idx = torch.tensor([1, 3, 5, 7]).cuda()
    data_2d = torch.zeros((32, 64)).cuda()
    data_3d = torch.zeros((32, 64, 128)).cuda()

    case_a = run_case("A: 1-D values, runtime-length index", OneDimValues(), data_2d, idx, True)
    case_b = run_case(
        "B: full-rank values, runtime-length index", FullRankValues(), data_3d, idx, True
    )
    case_c = run_case(
        "C: full-rank values, static-length index", FullRankValues(), data_3d, idx, False
    )

    reproduced = not case_a or not case_b
    print(f"\nstatic-extent control passed: {case_c}")
    print(f"reproduced: {reproduced}")
    return 0 if reproduced else 1


if __name__ == "__main__":
    sys.exit(main(argv=sys.argv))

output

torch          2.13.0a0+9186a08b2c.nv26.07
torch_tensorrt 2.14.0a0

===== A: 1-D values, runtime-length index =====
WARNING:torch_tensorrt.dynamo.partitioning.common:Dynamic input arg1_1 (shape: torch.Size([s95])) has no max bound for dim 0, attempting to use a sane default (max: min(1) * 2^12). Please set an upper bound using torch._dynamo.mark_dynamic or torch.export.Dim
ERROR:torch_tensorrt [TensorRT Conversion Context]:IBuilder::buildEngineWithConfig: Error Code 4: Internal Error (kOPT values for profile 0 violate shape constraints: [SLICE]-[aten_ops.index_put.default]-[index_put_expand_values]: ISliceLayer has out of bounds access on axis 1 Condition '<' violated: 31 >= 4. In operator() at /_src/runtime/core/shapeMachine.cpp:628 In evaluateShapeChecks at /_src/optimizer/common/shapeCompiler.cpp:1375)
CRITICAL:torch_tensorrt.dynamo.backend.backends:Halting compilation on build failure since pass_through_build_failures was specified as True. To return the default Torch implementation and avoid halting compilation on engine build failures, specify pass_through_build_failures=False.
Traceback (most recent call last):
  File "/w/repro.py", line 97, in run_case
    out = optimized(data, idx)
          ^^^^^^^^^^^^^^^^^^^^
  [...]
  File "/usr/local/lib/python3.12/dist-packages/torch_tensorrt/dynamo/conversion/_conversion.py", line 280, in interpret_module_to_result
    interpreter_result = interpreter.run()
                         ^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py", line 653, in run
    assert cuda_engine
           ^^^^^^^^^^^
torch._dynamo.exc.BackendCompilerFailed: backend='tensorrt' raised:
AssertionError: 

Set TORCHDYNAMO_VERBOSE=1 for the internal stack trace (please do this especially if you're reporting a bug to PyTorch). For even more developer context, set TORCH_LOGS="+dynamo"

----- A: 1-D values, runtime-length index: FAILED

===== B: full-rank values, runtime-length index =====
WARNING:torch_tensorrt.dynamo.partitioning.common:Dynamic input arg1_1 (shape: torch.Size([s95])) has no max bound for dim 0, attempting to use a sane default (max: min(1) * 2^12). Please set an upper bound using torch._dynamo.mark_dynamic or torch.export.Dim
----- B: full-rank values, runtime-length index: OK (32, 64, 128), matches eager

===== C: full-rank values, static-length index =====
----- C: full-rank values, static-length index: OK (32, 64, 128), matches eager

static-extent control passed: True
reproduced: True

Expected behavior

data[:, idx] = values with a rank-deficient values and a runtime-length index should
either be converted correctly -- broadcasting values across the free dimension the same way
Torch does -- or be declined.

Broadcasting across a dimension whose extent is only known at runtime is genuinely awkward to
express in ScatterND, since the whole index-materialisation strategy assumes the extents are
build-time constants. If it cannot be expressed, the right outcome is for the converter (or
its capability validator) to DECLINE the node, so the partitioner leaves aten.index_put in
Torch and the rest of the model still compiles -- rather than emitting a network TensorRT will
reject and taking the whole compilation down. The specific thing to fix in either case is line
1065: a pair like (32, -1) is a mismatch that happens to involve a dynamic extent, and
collapsing it to -1 turns a detectable problem into a malformed network.

Separate, smaller request: the diagnostics here are poor. The TensorRT build error is
logged, but what Python actually raises is a bare AssertionError with an empty message,
from assert cuda_engine at _TRTInterpreter.py:653. A user who does not have the TensorRT
logger output in front of them gets BackendCompilerFailed: backend='tensorrt' raised: AssertionError: and nothing else -- no layer name, no shape constraint, no mention of
index_put. Raising an exception that carries the builder's error text would make this class
of failure far cheaper to diagnose, independently of the broadcast bug.

Environment

Build information about Torch-TensorRT can be found by turning on debug messages

• Pytorch NGC container : 26.07-py3