[CI Fix] Propagate dtype in TorchAOBaseTensor._to_copy

test/prototype/mx_formats/test_mx_tensor.py

@@ -670,8 +670,9 @@ def test_index_select():
     reason="float8 in triton requires CUDA capability 8.9 or greater",
 )
 def test_cast_to_float8_e4m3fn_saturation_behavior():
-    # TODO(#1912): make the saturated cast work in eager mode and remove this
-    # test
+    # PyTorch now implements saturated casting to float8_e4m3fn in both eager
+    # and compiled modes. This test verifies the saturation behavior.
+    # Relates to issue #1912.
     max_val = torch.finfo(torch.float8_e4m3fn).max
 
     # create example data inside the representable range
@@ -694,13 +695,20 @@ def test_cast_to_float8_e4m3fn_saturation_behavior():
         device="cuda",
     )
 
-    # verify that in eager mode PyTorch casting to float8 is unsaturated
+    # verify that in eager mode PyTorch casting to float8 is saturated (no NaN)
     data_in_range_f8 = data_in_range_bf16.to(torch.float8_e4m3fn)
     data_out_of_range_f8 = data_out_of_range_bf16.to(torch.float8_e4m3fn)
     assert not torch.any(torch.isnan(data_in_range_f8))
-    assert torch.all(torch.isnan(data_out_of_range_f8))
+    assert not torch.any(torch.isnan(data_out_of_range_f8))
+    # Out-of-range values should be clamped to max_val
+    torch.testing.assert_close(
+        data_out_of_range_f8.to(torch.bfloat16),
+        torch.tensor([max_val, -1 * max_val], dtype=torch.bfloat16, device="cuda"),
+        atol=0,
+        rtol=0,
+    )
 
-    # verify that in triton, casting to float8 is saturated
+    # verify that in triton/compiled mode, casting to float8 is also saturated
     # for simplicity, use torch.compile to generate triton code
     def to_f8(x):
         x = x.to(torch.float8_e4m3fn)
@@ -711,8 +719,10 @@ def to_f8(x):
     data_out_of_range_f8_c = to_f8_c(data_out_of_range_bf16)
     assert not torch.any(torch.isnan(data_in_range_f8_c))
     assert not torch.any(torch.isnan(data_out_of_range_f8_c))
+    # Eager and compiled should produce the same saturated results
+    torch.testing.assert_close(data_in_range_f8, data_in_range_f8_c, atol=0, rtol=0)
     torch.testing.assert_close(
-        data_in_range_f8_c, data_out_of_range_f8_c, atol=0, rtol=0
+        data_out_of_range_f8, data_out_of_range_f8_c, atol=0, rtol=0
     )
 
 

test/test_utils.py

@@ -89,6 +89,46 @@ def __init__(self, qdata, attr="attr", device=None):
         kwargs = t._get_to_kwargs(device="cpu")
         self.assertFalse(kwargs["non_blocking"])
 
+    def test_to_copy_propagates_dtype_and_non_blocking(self):
+        """Verify that .to() propagates device, dtype, and non_blocking to inner tensors."""
+
+        class MyTensor(TorchAOBaseTensor):
+            tensor_data_names = ["qdata", "scale"]
+            tensor_attribute_names = ["device"]
+
+            def __new__(cls, qdata, scale, device=None):
+                if device is None:
+                    device = qdata.device
+                kwargs = {"device": device, "dtype": qdata.dtype}
+                r = torch.Tensor._make_wrapper_subclass(cls, qdata.shape, **kwargs)
+                r.qdata = qdata
+                r.scale = scale
+                return r
+
+            def __init__(self, qdata, scale, device=None):
+                pass
+
+        # Test dtype change propagates to inner tensors
+        t = MyTensor(
+            torch.randn(4, 4, dtype=torch.float32),
+            torch.tensor(1.0, dtype=torch.float32),
+        )
+        result = t.to(dtype=torch.float16)
+        self.assertEqual(result.dtype, torch.float16)
+        self.assertEqual(result.qdata.dtype, torch.float16)
+        self.assertEqual(result.scale.dtype, torch.float16)
+
+        # Test combined device and dtype change with non_blocking
+        t = MyTensor(
+            torch.randn(4, 4, dtype=torch.float32),
+            torch.tensor(1.0, dtype=torch.float32),
+        )
+        result = t.to(device="cpu", dtype=torch.bfloat16, non_blocking=True)
+        self.assertEqual(result.dtype, torch.bfloat16)
+        self.assertEqual(result.qdata.dtype, torch.bfloat16)
+        self.assertEqual(result.scale.dtype, torch.bfloat16)
+        self.assertEqual(result.device.type, "cpu")
+
     def _test_default_impls_helper(self, lp_tensor, lp_tensor_for_copy):
         # get `all_tensor_data_names` and `all_tensor_attribute_names`
         all_tensor_data_names = lp_tensor.tensor_data_names.copy()

torchao/utils.py

@@ -594,9 +594,12 @@ def _(func, types, args, kwargs):
         ):
             kwargs = self._get_to_kwargs(*args[1:], **kwargs)
             device = kwargs.pop("device")
+            dtype = kwargs.pop("dtype")
             non_blocking = kwargs.pop("non_blocking", False)
             tensors = [
-                getattr(self, name).to(device, non_blocking=non_blocking)
+                getattr(self, name).to(
+                    device=device, dtype=dtype, non_blocking=non_blocking
+                )
                 for name in self.tensor_data_names
             ]
             optional_tensors = []
@@ -605,7 +608,9 @@ def _(func, types, args, kwargs):
                     maybe_tensor = getattr(self, tensor_data_name)
                     if maybe_tensor is not None:
                         optional_tensors.append(
-                            maybe_tensor.to(device, non_blocking=non_blocking)
+                            maybe_tensor.to(
+                                device=device, dtype=dtype, non_blocking=non_blocking
+                            )
                         )
                     else:
                         optional_tensors.append(None)