🐛 [Bug] a SymInt engine output is returned as a 0-d tensor and never converted back for scalar consumers

[SrivastavaKshitij]

Bug Description

When a torch.SymInt is an output of a TensorRT engine, it comes back as a 0-d tensor and
nothing converts it back to a scalar before the Torch consumer that was traced against a
scalar. The consumer then rejects it.

Two pieces combine:

1. py/torch_tensorrt/dynamo/utils.py :: get_output_dtypes deliberately types a SymInt
   output as int64:

            elif isinstance(output_meta, torch.SymInt):
                output_dtypes.append(dtype.int64)

2. tensorrt::execute_engine can only return tensors. Its declaration, as printed by the
   error below, is:

   tensorrt::execute_engine(Tensor[] input_tensors, __torch__.torch.classes.tensorrt.Engine engine) -> Tensor[]
   

So py/torch_tensorrt/dynamo/_exporter.py :: inline_trt_modules replaces the engine
call_module with an execute_engine call plus operator.getitem nodes, and wires the
getitem for the SymInt output straight into the consumer:

                trt_node = gm.graph.call_function(
                    torch.ops.tensorrt.execute_engine.default,
                    (trt_module_node.args, engine_node),
                )

No aten.item / _local_scalar_dense is inserted on that edge. The consumer -- here
aten.arange.start_step, whose end argument is a Scalar -- receives a 0-d int64 tensor
and fails.

The information needed to do the right thing is present and even recorded: the engine's
serialized metadata marks that output is_scalar: True (see the output quoted below). But
_apply_symbolic_shape_expressions in runtime/meta_ops/register_meta_ops.py builds
torch.empty(...) for every entry in outputs regardless of the flag, and the op schema could
not express a scalar return even if it read it.

Where the failure is reported is misleading. The traceback names aten::arange in the
generated forward and does not mention torch_tensorrt at all, because by then the exporter
has finished its work and left the scene. The provenance is only visible in the artifacts the
exporter produced, which the script prints (engine metadata and the inlined graph).

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
from torch._subclasses.fake_tensor import FakeTensorMode
from torch.fx.experimental.symbolic_shapes import ShapeEnv
from torch_tensorrt.dynamo._exporter import transform

ROWS = 16


class SymIntEngineOutput(torch.nn.Module):
    """A symbolic dimension read inside the engine and consumed by Torch outside it."""

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        y = x * 2.0
        n = y.shape[0]
        return y.sum() + torch.arange(0, n, device=x.device).float().sum()


def build_engine() -> torch.fx.GraphModule:
    """Compiles the model with `arange` pinned to Torch so the SymInt leaves the engine."""
    model = SymIntEngineOutput().eval().cuda()
    x = torch.randn((ROWS,), device="cuda")
    rows = torch.export.Dim("rows", min=2, max=64)
    exported = torch.export.export(model, (x,), dynamic_shapes={"x": {0: rows}})
    return torch_tensorrt.dynamo.compile(
        exported,
        inputs=(x,),
        min_block_size=1,
        pass_through_build_failures=True,
        torch_executed_ops={"torch.ops.aten.arange.start_step"},
    )


def main(argv: list[str] | tuple[str, ...] = ()) -> int:
    """Inlines the engine call and runs it, so `arange` receives the engine's 0-d tensor."""
    del argv
    print(f"torch          {torch.__version__}")
    print(f"torch_tensorrt {torch_tensorrt.__version__}")

    trt_gm = build_engine()
    print(trt_gm.graph, flush=True)
    for name, child in trt_gm.named_children():
        exprs = getattr(child, "symbolic_shape_expressions", None)
        if exprs is not None:
            print(f"{name} inputs : {exprs.get('inputs')}")
            print(f"{name} outputs: {exprs.get('outputs')}", flush=True)

    # transform() is what torch_tensorrt.dynamo.export()/AOTInductor consume: it replaces
    # the TRT submodule with a direct execute_engine call whose results are Tensors.
    inlined = transform(trt_gm)
    inlined.recompile()
    print(inlined.graph, flush=True)

    reproduced = False
    with FakeTensorMode(shape_env=ShapeEnv()):
        x_fake = torch.empty((ROWS,), dtype=torch.float32, device="cuda")
        try:
            out = inlined(x_fake)
            print(f"inlined graph returned {out}")
        except Exception as exc:  # pylint: disable=broad-except
            traceback.print_exc()
            reproduced = "for argument 'end'" in str(exc) and "aten::arange" in str(exc)

    print(f"\nreproduced: {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
graph():
    %x : [num_users=1] = placeholder[target=x]
    %_run_on_acc_0 : [num_users=2] = call_module[target=_run_on_acc_0](args = (%x,), kwargs = {})
    %getitem : [num_users=1] = call_function[target=operator.getitem](args = (%_run_on_acc_0, 0), kwargs = {})
    %getitem_1 : [num_users=1] = call_function[target=operator.getitem](args = (%_run_on_acc_0, 1), kwargs = {})
    %_run_on_gpu_1 : [num_users=1] = call_module[target=_run_on_gpu_1](args = (%getitem,), kwargs = {})
    %_run_on_acc_2 : [num_users=1] = call_module[target=_run_on_acc_2](args = (%_run_on_gpu_1, %getitem_1), kwargs = {})
    return (_run_on_acc_2,)
_run_on_acc_0 inputs : [{'shape_exprs': [s77], 'dtype': torch.float32, 'name': 'x'}]
_run_on_acc_0 outputs: [{'shape_exprs': [], 'dtype': torch.int64, 'is_scalar': True}, {'shape_exprs': [], 'dtype': torch.float32}]
_run_on_acc_2 inputs : [{'shape_exprs': [s77], 'dtype': torch.int64, 'name': 'arange'}, {'shape_exprs': [], 'dtype': torch.float32, 'name': 'sum_1'}]
_run_on_acc_2 outputs: [{'shape_exprs': [], 'dtype': torch.float32}]
WARNING:py.warnings:/usr/local/lib/python3.12/dist-packages/torch_tensorrt/dynamo/_exporter.py:500: UserWarning: Attempted to insert a get_attr Node with no underlying reference in the owning GraphModule! Call GraphModule.add_submodule to add the necessary submodule, GraphModule.add_parameter to add the necessary Parameter, or nn.Module.register_buffer to add the necessary buffer
  engine_node = gm.graph.get_attr(engine_name)

graph():
    %x : [num_users=1] = placeholder[target=x]
    %_run_on_acc_0_engine : [num_users=1] = get_attr[target=_run_on_acc_0_engine]
    %execute_engine_default : [num_users=2] = call_function[target=torch.ops.tensorrt.execute_engine.default](args = ((%x,), %_run_on_acc_0_engine), kwargs = {})
    %getitem : [num_users=1] = call_function[target=operator.getitem](args = (%execute_engine_default, 0), kwargs = {})
    %getitem_1 : [num_users=1] = call_function[target=operator.getitem](args = (%execute_engine_default, 1), kwargs = {})
    %arange : [num_users=1] = call_function[target=torch.ops.aten.arange.start_step](args = (0, %getitem), kwargs = {layout: torch.strided, device: cuda:0, pin_memory: False})
    %_run_on_acc_2_engine : [num_users=1] = get_attr[target=_run_on_acc_2_engine]
    %execute_engine_default_1 : [num_users=1] = call_function[target=torch.ops.tensorrt.execute_engine.default](args = ((%arange, %getitem_1), %_run_on_acc_2_engine), kwargs = {})
    %getitem_2 : [num_users=1] = call_function[target=operator.getitem](args = (%execute_engine_default_1, 0), kwargs = {})
    return (getitem_2,)
Traceback (most recent call last):
  File "/w/repro.py", line 92, in main
    out = inlined(x_fake)
          ^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/graph_module.py", line 1000, in call_wrapped
    return self._wrapped_call(self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[...]
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1789, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<eval_with_key>.38", line 10, in forward
    arange = torch.ops.aten.arange.start_step(0, getitem, layout = torch.strided, device = device(type='cuda', index=0), pin_memory = False);  getitem = None
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_ops.py", line 875, in __call__
    return self._op(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: aten::arange() Expected a value of type 'number' for argument 'end' but instead found type 'FakeTensor'.
Position: 1
Value: FakeTensor(..., device='cuda:0', size=(), dtype=torch.int64)
Declaration: aten::arange.start_step(Scalar start, Scalar end, Scalar step=1, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor
Cast error details: Cannot cast FakeTensor(..., device='cuda:0', size=(), dtype=torch.int64) to number

reproduced: True

Expected behavior

When an engine output was a SymInt in the traced graph, inline_trt_modules should restore a
scalar on that edge -- i.e. insert an aten._local_scalar_dense (or equivalent .item()) node
between the getitem and its consumers, so consumers traced against a Scalar keep receiving
one. The is_scalar: True flag already recorded in the engine's metadata (and the SymInt
branch in get_output_dtypes that produced it) identifies exactly which outputs need it.

The engine meta kernel should agree with that: _apply_symbolic_shape_expressions currently
builds a torch.empty(...) for every entry in outputs and should instead produce a symbolic
scalar for entries flagged is_scalar, so the fake and real paths report the same types.

If restoring the scalar is not desired, the alternative is for the partitioner to decline to
put a SymInt-producing node at an engine boundary in the first place, so the whole
computation stays on one side. What should not happen is a successfully "compiled" module that
raises a schema RuntimeError the first time it is run.

Environment

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

• Pytorch NGC container : 26.07-py3