Fix/converter dtype promotion

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
 

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)
 

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
     )

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

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

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