From 91b21b9d0606d5e7216ca49fe099a85e90ef7479 Mon Sep 17 00:00:00 2001 From: Kshitij Srivastava Date: Fri, 28 Aug 2026 02:48:09 +0000 Subject: [PATCH] fix: unify LINSPACE alpha/beta dtype in the arange converter `arange` resolves a common `value_dtype` for the sequence and passes it to each `get_trt_tensor` call. `get_trt_tensor` applies that dtype when it constructs a constant, and returns a value that is already an ITensor unchanged: elif isinstance(input_val, TRTTensor): return input_val So a dynamic `start` keeps the dtype of its incoming ITensor, while a literal `step` becomes a constant of `value_dtype`. TensorRT asks that the LINSPACE `alpha` (input 1) and `beta` (input 2) have the same type, and reports: IFillLayer `alpha` and `beta` must have the same type. `alpha` is of type Int32 but `beta` is of type Int64. This change routes every operand through a small helper that casts after `get_trt_tensor`, so the resolved dtype holds for all of them. `cast_trt_tensor` returns the tensor unchanged when the dtype already matches, so paths that were already consistent are unaffected. The tests cover three cases with a dynamic bound: an int32 `start`, an int64 `start`, and both bounds dynamic. Of those, only the int64 `start` works today, because the sequence dtype for integer operands is already int64 and matches. The other two pass with this change, and the full arange converter suite stays green. --- .../dynamo/conversion/impl/arange.py | 34 +++++++-- .../py/dynamo/conversion/test_arange_aten.py | 72 +++++++++++++++++++ 2 files changed, 100 insertions(+), 6 deletions(-) diff --git a/py/torch_tensorrt/dynamo/conversion/impl/arange.py b/py/torch_tensorrt/dynamo/conversion/impl/arange.py index 81fbbae590..bce952ae8b 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/arange.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/arange.py @@ -41,6 +41,28 @@ def _sequence_dtype( return trt.DataType.INT64 +def _operand_as( + ctx: ConversionContext, + target: Target, + source_ir: Optional[SourceIR], + name: str, + value: Union[int, float, TRTTensor], + dtype: trt.DataType, + min_rank: int, +) -> TRTTensor: + """ + Materialize an arange operand as an ITensor of exactly `dtype`. + + ``get_trt_tensor`` applies its ``dtype`` argument when it constructs a constant, and + returns a value that is already an ITensor unchanged. Casting afterwards lets the + resolved dtype hold for every operand, which is what LINSPACE asks of ``alpha`` and + ``beta``. ``cast_trt_tensor`` returns the tensor unchanged when the dtype already + matches. + """ + operand = get_trt_tensor(ctx, value, name, dtype, min_rank=min_rank) + return cast_trt_tensor(ctx, operand, dtype, name + "_casted", target, source_ir) + + def arange( ctx: ConversionContext, target: Target, @@ -62,8 +84,8 @@ def arange( # If any argument is a TRT tensor, use dynamic arange with a Fill layer if any(isinstance(x, TRTTensor) for x in (start, end, step)): value_dtype = _sequence_dtype(dtype, start, end, step) - start_rank_0 = get_trt_tensor( - ctx, start, name + "_start_rank_0", value_dtype, min_rank=0 + start_rank_0 = _operand_as( + ctx, target, source_ir, name + "_start_rank_0", start, value_dtype, 0 ) # LINSPACE's start input requires rank 0; if the upstream ITensor came in # as rank-1 (e.g. a SymInt materialized by a sym_size op), reshape it. @@ -75,11 +97,11 @@ def arange( ) start_rank_0 = squeeze_layer.get_output(0) - start_rank_1 = get_trt_tensor( - ctx, start, name + "_start_rank_1", value_dtype, min_rank=1 + start_rank_1 = _operand_as( + ctx, target, source_ir, name + "_start_rank_1", start, value_dtype, 1 ) - end = get_trt_tensor(ctx, end, name + "_end", value_dtype, min_rank=1) - step = get_trt_tensor(ctx, step, name + "_step", value_dtype, min_rank=1) + end = _operand_as(ctx, target, source_ir, name + "_end", end, value_dtype, 1) + step = _operand_as(ctx, target, source_ir, name + "_step", step, value_dtype, 1) # The number of elements is ceil((end - start) / step), computed as # -floor((start - end) / step) so that the whole expression stays in the diff --git a/tests/py/dynamo/conversion/test_arange_aten.py b/tests/py/dynamo/conversion/test_arange_aten.py index e48f14b93e..c3c70acc75 100644 --- a/tests/py/dynamo/conversion/test_arange_aten.py +++ b/tests/py/dynamo/conversion/test_arange_aten.py @@ -93,6 +93,78 @@ def forward(self, end_tensor): use_dynamo_tracer=False, ) + @parameterized.expand([("int32", torch.int32), ("int64", torch.int64)]) + def test_arange_dynamic_start(self, _, dtype): + """A dynamic `start` reaches the Fill layer as an ITensor. + + LINSPACE requires `alpha` (start) and `beta` (step) to share a dtype. A dynamic + `start` keeps the dtype of its incoming ITensor, while a literal `step` is + materialized as a constant of the sequence dtype, so the two can differ. + + Three cases with a dynamic bound are covered here: an int32 `start`, an int64 + `start`, and both bounds dynamic. Of those, only the int64 `start` works today, + because the sequence dtype for integer operands is already int64 and matches. + """ + + class Arange(nn.Module): + def forward(self, start_tensor): + return torch.ops.aten.arange.start_step(start_tensor, 10, 1) + + pyt_input = 2 + inputs = [ + torch_tensorrt.Input( + min_shape=(0,), + opt_shape=(2,), + max_shape=(5,), + dtype=dtype, + torch_tensor=torch.tensor(pyt_input, dtype=dtype).cuda(), + is_shape_tensor=True, + ) + ] + self.run_test_with_dynamic_shape( + Arange(), + inputs, + use_example_tensors=False, + check_dtype=False, + pyt_inputs=[pyt_input], + use_dynamo_tracer=False, + ) + + def test_arange_dynamic_start_and_end(self): + """Both bounds dynamic, so neither is a constant carrying the sequence dtype.""" + + class Arange(nn.Module): + def forward(self, start_tensor, end_tensor): + return torch.ops.aten.arange.start_step(start_tensor, end_tensor, 1) + + pyt_inputs = [2, 9] + inputs = [ + torch_tensorrt.Input( + min_shape=(0,), + opt_shape=(2,), + max_shape=(5,), + dtype=torch.int32, + torch_tensor=torch.tensor(pyt_inputs[0], dtype=torch.int32).cuda(), + is_shape_tensor=True, + ), + torch_tensorrt.Input( + min_shape=(6,), + opt_shape=(9,), + max_shape=(12,), + dtype=torch.int64, + torch_tensor=torch.tensor(pyt_inputs[1], dtype=torch.int64).cuda(), + is_shape_tensor=True, + ), + ] + self.run_test_with_dynamic_shape( + Arange(), + inputs, + use_example_tensors=False, + check_dtype=False, + pyt_inputs=pyt_inputs, + use_dynamo_tracer=False, + ) + if __name__ == "__main__": run_tests()