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 ) 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()