🐛 [Bug] rank-8+ tensors cross the partition boundary and fail in add_input, not at their producer

Bug Description

trt.Dims holds at most trt.Dims.MAX_DIMS axes, which is 8 (verified in this container:
python -c "import tensorrt as trt; print(trt.Dims.MAX_DIMS)" prints 8). Partitioning only
asks whether a converter exists for a node:

• py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py :: OpSupportTester.is_node_supported
• py/torch_tensorrt/dynamo/partitioning/_global_partitioner.py :: TorchTensorRTOperatorSupport.is_node_supported

Neither asks whether the tensors crossing the partition boundary are representable in
trt.Dims. The crux is that declining only the producing op is not enough. In the
reproducer, a custom op with no converter produces a rank-10 tensor, so the partitioner
correctly declines it and it runs in Torch. But its consumer (* 2.0) does have a converter,
so the consumer is accepted into a TensorRT subgraph, and that drags the rank-10 tensor across
the boundary as a network input. It then fails as an input, not at its producer, in

# py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py :: placeholder
return self.ctx.net.add_input(
    name=target,
    shape=tuple(shape),
    dtype=trt_input_dtype,
)

with a pybind overload TypeError (the 10-element tuple cannot be converted to
tensorrt.tensorrt.Dims). Because the failing node is a placeholder, the error carries no
useful provenance - the report ends with Original traceback: None, and nothing points at the
op that actually produced the over-rank tensor.

To Reproduce

Steps to reproduce the behavior:

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 traceback

import torch

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

COMPILE_OPTIONS = {"pass_through_build_failures": True, "min_block_size": 1}
RANK_10_SHAPE = (1, 2, 1, 1, 1, 1, 1, 1, 4, 8)
RANK_8_SHAPE = (1, 2, 1, 1, 1, 1, 4, 8)


@torch.library.custom_op("repro::high_rank", mutates_args=())
def high_rank(x: torch.Tensor) -> torch.Tensor:
    """Reshapes to rank 10. Opaque to TensorRT, so the partitioner has to decline it."""
    return x.reshape(RANK_10_SHAPE).clone()


@high_rank.register_fake
def _high_rank_fake(x: torch.Tensor) -> torch.Tensor:
    return x.new_empty(RANK_10_SHAPE)


@torch.library.custom_op("repro::low_rank", mutates_args=())
def low_rank(x: torch.Tensor) -> torch.Tensor:
    """Control counterpart of `high_rank`: same structure, rank 8."""
    return x.reshape(RANK_8_SHAPE).clone()


@low_rank.register_fake
def _low_rank_fake(x: torch.Tensor) -> torch.Tensor:
    return x.new_empty(RANK_8_SHAPE)


class HighRankBoundary(torch.nn.Module):
    """The rank-10 producer falls back to Torch; the consumer is converted."""

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.ops.repro.high_rank(x) * 2.0


class LowRankBoundary(torch.nn.Module):
    """Control: identical structure, rank 8, within trt.Dims.MAX_DIMS."""

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.ops.repro.low_rank(x) * 2.0


def compile_and_run(model: torch.nn.Module, args: tuple[torch.Tensor, ...]) -> BaseException | None:
    """Compiles with the TensorRT backend and runs it. Returns the exception, or None."""
    torch._dynamo.reset()
    optimized = torch.compile(model, backend="tensorrt", options=COMPILE_OPTIONS)
    try:
        optimized(*args)
    except BaseException as exc:  # noqa: BLE001  # a bug demonstration: report, do not handle
        traceback.print_exc()
        return exc
    return None


def main() -> int:
    print(f"torch {torch.__version__}")
    print(f"torch_tensorrt {torch_tensorrt.__version__}")

    x = torch.rand(2, 4, 8, device="cuda")

    print("\n=== control: rank-8 tensor across the partition boundary ===")
    control_exc = compile_and_run(LowRankBoundary().cuda().eval(), (x,))
    control_ok = control_exc is None
    print(f"control compiled: {control_ok}")

    print("\n=== bug case: rank-10 tensor across the partition boundary ===")
    bug_exc = compile_and_run(HighRankBoundary().cuda().eval(), (x,))
    if bug_exc is None:
        print("bug case compiled cleanly -- the bug did not reproduce")
        reproduced = False
    else:
        tb = "".join(traceback.format_exception(type(bug_exc), bug_exc, bug_exc.__traceback__))
        names_add_input = "add_input" in str(bug_exc) or "add_input" in tb
        in_placeholder = "_TRTInterpreter.py" in tb and ", in placeholder" in tb
        print(f"failure names add_input                 : {names_add_input}")
        print(f"traceback names placeholder             : {in_placeholder}")
        print(f"exception: {type(bug_exc).__name__}: {str(bug_exc)[:400]}")
        reproduced = names_add_input and in_placeholder

    print(f"\ncontrol compiles: {control_ok}")
    print(f"reproduced: {reproduced}")
    return 0 if (reproduced and control_ok) else 1


if __name__ == "__main__":
    raise SystemExit(main())

output

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

=== control: rank-8 tensor across the partition boundary ===
control compiled: True

=== bug case: rank-10 tensor across the partition boundary ===
Traceback (most recent call last):
  File "/w/repro.py", line 88, in compile_and_run
    optimized(*args)
[...]
  File "/usr/local/lib/python3.12/dist-packages/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py", line 384, in _construct_trt_network_def
    super().run()
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/interpreter.py", line 197, in run
    self.env[node] = self.run_node(node)
                     ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py", line 686, in run_node
    trt_node: torch.fx.Node = super().run_node(n)
                              ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/interpreter.py", line 294, in run_node
    return getattr(self, n.op)(n.target, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py", line 749, in placeholder
    return self.ctx.net.add_input(
           ^^^^^^^^^^^^^^^^^^^^^^^
torch._dynamo.exc.BackendCompilerFailed: backend='tensorrt' raised:
TypeError: add_input(): incompatible function arguments. The following argument types are supported:
    1. (self: tensorrt.tensorrt.INetworkDefinition, name: str, dtype: tensorrt.tensorrt.DataType, shape: tensorrt.tensorrt.Dims) -> tensorrt.tensorrt.ITensor

Invoked with: <tensorrt.tensorrt.INetworkDefinition object at 0x727bfad294b0>; kwargs: name='high_rank', shape=(1, 2, 1, 1, 1, 1, 1, 1, 4, 8), dtype=<DataType.FLOAT: 0>

While executing %high_rank : [num_users=1] = placeholder[target=high_rank]
Original traceback:
None

failure names add_input                 : True
traceback names placeholder             : True

control compiles: True
reproduced: True

Expected behavior

The partitioner should decline any node whose output or any of whose inputs has rank greater
than trt.Dims.MAX_DIMS, so the whole affected region falls back to Torch and compilation of
the rest of the model succeeds. Aborting the entire torch.compile because one boundary tensor
has 10 axes is the wrong outcome: rank > 8 is simply not expressible in TensorRT, and that is
exactly the situation partitioning exists to handle.

There is already a precedent for the "consider a node's inputs too" pattern in the same place.
TorchTensorRTOperatorSupport._has_complex_dtype (added in #4119) is called from both
is_node_supported implementations and does exactly this:

if _dtype(node) in COMPLEX_DTYPES:
    return True
for arg in node.all_input_nodes:
    if _dtype(arg) in COMPLEX_DTYPES:
        return True
return False

A rank check belongs alongside it - same call site, same shape of predicate, using
trt.Dims.MAX_DIMS instead of COMPLEX_DTYPES. We have not implemented or tested such a
patch, so treat that as a suggestion rather than a verified fix.

Separately: even if the decline is not added, the diagnostics deserve improvement. A pybind
overload-resolution TypeError listing the accepted C++ signatures, attached to a placeholder
with Original traceback: None, gives a user nothing to act on. An explicit check in
placeholder (or better, at partitioning time) with a message naming the tensor, its rank and
the limit would be far more actionable.

Environment

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

• Pytorch NGC container : 26.07-py3

Start with OpSupportTester.is_node_supported in py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py and TorchTensorRTOperatorSupport.is_node_supported in _global_partitioner.py, then compare their use of _has_complex_dtype. Use the supplied repro.py as the regression case, checking both rank-8 and rank-10 boundaries. Done means unsupported high-rank boundary regions fall back to Torch without the placeholder add_input failure.