From fcdfe12367fc1632f5011ca05b3ffbbf02245950 Mon Sep 17 00:00:00 2001 From: Dongwoon Hyun Date: Tue, 18 Aug 2026 15:40:03 -0700 Subject: [PATCH 1/3] fix: correct dtype promotion bugs in converter layer Three related fixes for use_explicit_typing correctness: 1. scatter: replace `np.ones` with `np.full` to avoid float64 scalar promotion, and cast src_tensor dtype when it mismatches input (e.g. after argmax/topk emit INT32 instead of INT64). 2. promote_trt_tensors_to_same_dtype: preserve bool dtype when both operands are bool, rather than erroneously promoting to int32. 3. convert_binary_elementwise: use torch.result_type for scalar-tensor dtype promotion instead of unconditionally casting the scalar to the tensor's dtype (e.g. int64_tensor * float_scalar now correctly yields float32, not int64). --- .../dynamo/conversion/converter_utils.py | 5 ++++- .../dynamo/conversion/impl/elementwise/base.py | 16 +++++++++++----- .../dynamo/conversion/impl/select.py | 7 ++++++- 3 files changed, 21 insertions(+), 7 deletions(-) diff --git a/py/torch_tensorrt/dynamo/conversion/converter_utils.py b/py/torch_tensorrt/dynamo/conversion/converter_utils.py index 816f8aec58..7ceabc2af7 100644 --- a/py/torch_tensorrt/dynamo/conversion/converter_utils.py +++ b/py/torch_tensorrt/dynamo/conversion/converter_utils.py @@ -1180,8 +1180,11 @@ def promote_trt_tensors_to_same_dtype( promoted_dtype = trt.float32 else: promoted_dtype = trt.float16 + elif lhs.dtype == trt.bool and rhs.dtype == trt.bool: + # Case 2: If both tensors are bool types, preserve bool + promoted_dtype = trt.bool else: - # Case 2: If both tensors are int types (e.g., int32, int64), promote to int32 + # Case 3: If both tensors are int types (e.g., int32, int64), promote to int32 # (Note: TensorRT does not support int64 for many ops like select/where) promoted_dtype = trt.int32 diff --git a/py/torch_tensorrt/dynamo/conversion/impl/elementwise/base.py b/py/torch_tensorrt/dynamo/conversion/impl/elementwise/base.py index 097a81b8d1..72276bd9f3 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/elementwise/base.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/elementwise/base.py @@ -123,13 +123,19 @@ def convert_binary_elementwise( # into [], meaning that the result will have shape [], which is what we # expect. # - # Note that the dtype here is supposed to be the same as the scalar - # dtype but we don't have a way to detect whether it makes sense for the - # scalar to be float or half. Hence we go with the lhs dtype. + # Use torch.result_type to determine the scalar's dtype per PyTorch's promotion rules if is_lhs_trt_tensor and isinstance(rhs_val, (float, int, bool)): - rhs_val = to_torch(rhs_val, dtype=lhs_dtype) + lhs_torch_dtype = _enums.dtype._from(lhs_dtype).to(torch.dtype) + _rhs_dtype = _enums.dtype._from( + torch.result_type(torch.empty(0, dtype=lhs_torch_dtype), rhs_val) + ).to(trt.DataType) + rhs_val = to_torch(rhs_val, dtype=_rhs_dtype) if is_rhs_trt_tensor and isinstance(lhs_val, (float, int, bool)): - lhs_val = to_torch(lhs_val, dtype=rhs_dtype) + rhs_torch_dtype = _enums.dtype._from(rhs_dtype).to(torch.dtype) + _lhs_dtype = _enums.dtype._from( + torch.result_type(torch.empty(0, dtype=rhs_torch_dtype), lhs_val) + ).to(trt.DataType) + lhs_val = to_torch(lhs_val, dtype=_lhs_dtype) lhs_val = get_trt_tensor(ctx, lhs_val, f"{name}_lhs", lhs_dtype) rhs_val = get_trt_tensor(ctx, rhs_val, f"{name}_rhs", rhs_dtype) diff --git a/py/torch_tensorrt/dynamo/conversion/impl/select.py b/py/torch_tensorrt/dynamo/conversion/impl/select.py index 7b8a77dd56..2838681c04 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/select.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/select.py @@ -518,7 +518,7 @@ def scatter( else: # Static shape: use numpy to create the filled tensor src_tensor = get_trt_tensor( - ctx, src * np.ones(index_shape_list), name + "_value_tensor" + ctx, np.full(index_shape_list, src), name + "_value_tensor" ) src_tensor = cast_trt_tensor( ctx, src_tensor, input.dtype, name + "_cast_value_tensor" @@ -527,6 +527,11 @@ def scatter( elif not (isinstance(src, TRTTensor)): src_tensor = get_trt_tensor(ctx, src, name + "_src_tensor") + if isinstance(src_tensor, TRTTensor) and src_tensor.dtype != input.dtype: + src_tensor = cast_trt_tensor( + ctx, src_tensor, input.dtype, name + "_cast_src_tensor" + ) + scatter_layer = ctx.net.add_scatter( input, index, src_tensor, trt.ScatterMode.ELEMENT ) From e75c4a10d31c478b80a45dfb858c2335cdd358f0 Mon Sep 17 00:00:00 2001 From: Dongwoon Hyun Date: Tue, 18 Aug 2026 16:23:29 -0700 Subject: [PATCH 2/3] tests: add regression tests for dtype promotion fixes --- .../dynamo/conversion/test_binary_ops_aten.py | 18 ++++++++++++++++ .../py/dynamo/conversion/test_scatter_aten.py | 21 +++++++++++++++++++ tests/py/dynamo/conversion/test_where_aten.py | 12 +++++++++++ 3 files changed, 51 insertions(+) diff --git a/tests/py/dynamo/conversion/test_binary_ops_aten.py b/tests/py/dynamo/conversion/test_binary_ops_aten.py index 16b82b9858..df92d5ba98 100644 --- a/tests/py/dynamo/conversion/test_binary_ops_aten.py +++ b/tests/py/dynamo/conversion/test_binary_ops_aten.py @@ -257,5 +257,23 @@ def forward(self, x): self.run_test(m, inputs) +class TestScalarTensorDtypePromotion(DispatchTestCase): + def test_int_tensor_times_float_scalar_promotes_to_float(self): + class TestModule(nn.Module): + def forward(self, x): + return x * 8.2 + + inputs = [torch.tensor([1, 2, 3, 4], dtype=torch.int64)] + self.run_test(TestModule(), inputs) + + def test_float_scalar_times_int_tensor_promotes_to_float(self): + class TestModule(nn.Module): + def forward(self, x): + return 8.2 * x + + inputs = [torch.tensor([1, 2, 3, 4], dtype=torch.int64)] + self.run_test(TestModule(), inputs) + + if __name__ == "__main__": run_tests() diff --git a/tests/py/dynamo/conversion/test_scatter_aten.py b/tests/py/dynamo/conversion/test_scatter_aten.py index bb4c2eb440..396d056c93 100644 --- a/tests/py/dynamo/conversion/test_scatter_aten.py +++ b/tests/py/dynamo/conversion/test_scatter_aten.py @@ -316,5 +316,26 @@ def forward(self, input, index, src): self.run_test_with_dynamic_shape(TestModule(), input_specs) +class TestScatterDtypeFixConverter(DispatchTestCase): + def test_scatter_value_bool_scalar_no_float64_promotion(self): + class TestModule(torch.nn.Module): + def forward(self, x): + index = torch.tensor([[0, 1], [1, 0]], dtype=torch.int64) + return torch.ops.aten.scatter.value(x, 1, index, True) + + inputs = [torch.zeros(2, 4, dtype=torch.float32)] + self.run_test(TestModule(), inputs) + + def test_scatter_src_dtype_mismatch_after_argmax(self): + class TestModule(torch.nn.Module): + def forward(self, x): + index = torch.tensor([[0, 1], [1, 0]], dtype=torch.int64) + src = torch.argmax(x, dim=1, keepdim=True).expand(2, 2) + return torch.ops.aten.scatter.src(x, 1, index, src) + + inputs = [torch.arange(8, dtype=torch.int64).reshape(2, 4)] + self.run_test(TestModule(), inputs) + + if __name__ == "__main__": run_tests() diff --git a/tests/py/dynamo/conversion/test_where_aten.py b/tests/py/dynamo/conversion/test_where_aten.py index d965889861..9649e88f58 100644 --- a/tests/py/dynamo/conversion/test_where_aten.py +++ b/tests/py/dynamo/conversion/test_where_aten.py @@ -146,5 +146,17 @@ def forward(self, condition, x, y): self.run_test_with_dynamic_shape(Where(), input_specs) +class TestWhereBoolDtypePreservation(DispatchTestCase): + def test_where_bool_bool_output_is_bool_not_int32(self): + class Where(nn.Module): + def forward(self, cond, a, b): + return torch.ops.aten.where.self(cond, a, b) + + cond = torch.tensor([True, False, True, False]) + a = torch.tensor([True, True, False, False]) + b = torch.tensor([False, False, True, True]) + self.run_test(Where(), [cond, a, b]) + + if __name__ == "__main__": run_tests() From 6353067d6c31bc8583e5b881e8e47f57419d433e Mon Sep 17 00:00:00 2001 From: Dongwoon Hyun Date: Fri, 28 Aug 2026 10:23:44 -0700 Subject: [PATCH 3/3] Derived dtype directly from input and added scalar test for float --- py/torch_tensorrt/dynamo/conversion/impl/select.py | 12 ++++++------ tests/py/dynamo/conversion/test_scatter_aten.py | 11 +++++++++-- 2 files changed, 15 insertions(+), 8 deletions(-) diff --git a/py/torch_tensorrt/dynamo/conversion/impl/select.py b/py/torch_tensorrt/dynamo/conversion/impl/select.py index 2838681c04..e34af36eef 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/select.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/select.py @@ -20,7 +20,7 @@ ) from torch_tensorrt.dynamo.conversion.impl.elementwise import convert_binary_elementwise from torch_tensorrt.dynamo.conversion.impl.shape import shape as get_shape -from torch_tensorrt.dynamo.utils import DYNAMIC_DIM +from torch_tensorrt.dynamo.utils import DYNAMIC_DIM, Frameworks, unified_dtype_converter _LOGGER: logging.Logger = logging.getLogger(__name__) @@ -516,12 +516,12 @@ def scatter( input.dtype, ) else: - # Static shape: use numpy to create the filled tensor + # Static shape: use numpy to create the filled tensor. + input_np_dtype = unified_dtype_converter(input.dtype, Frameworks.NUMPY) src_tensor = get_trt_tensor( - ctx, np.full(index_shape_list, src), name + "_value_tensor" - ) - src_tensor = cast_trt_tensor( - ctx, src_tensor, input.dtype, name + "_cast_value_tensor" + ctx, + np.full(index_shape_list, src, dtype=input_np_dtype), + name + "_value_tensor", ) # scatter.src elif not (isinstance(src, TRTTensor)): diff --git a/tests/py/dynamo/conversion/test_scatter_aten.py b/tests/py/dynamo/conversion/test_scatter_aten.py index 396d056c93..24fd832249 100644 --- a/tests/py/dynamo/conversion/test_scatter_aten.py +++ b/tests/py/dynamo/conversion/test_scatter_aten.py @@ -63,6 +63,12 @@ def forward(self, input): torch.tensor([[0, 1, 2, 0], [1, 2, 1, 1]]), 1, ), + ( + "scatter_one_dim_float_scalar_value", + 1, + torch.tensor([[0, 1, 2, 0]]), + 2.5, + ), ] ) def test_scatter_index_input(self, _, dim, index, value): @@ -73,9 +79,10 @@ def __init__(self): def forward(self, input, index): return torch.ops.aten.scatter.value(input, dim, index, value) - input = torch.zeros(3, 5, dtype=torch.int32) + dtype = torch.float32 if isinstance(value, float) else torch.int32 + input = torch.zeros(3, 5, dtype=dtype) inputs = [input, index] - self.run_test(TestModule(), inputs, int32_reqd=True) + self.run_test(TestModule(), inputs, int32_reqd=(dtype == torch.int32)) class TestScatterSrcConverter(DispatchTestCase):