[Bugfix] Add MPS float64->float32 downcast (#3548)

Author: bsprenger

test/test_specs.py

@@ -4639,6 +4639,39 @@ def test_index_select_stacked_not_supported(self):
             torch.index_select(stacked_spec, dim=0, index=torch.tensor([0]))
 
 
+def _has_mps():
+    if hasattr(torch, "mps") and hasattr(torch.mps, "is_available"):
+        return torch.mps.is_available()
+    return (
+        getattr(torch.backends, "mps", None) is not None
+        and torch.backends.mps.is_available()
+    )
+
+
+@pytest.mark.skipif(not _has_mps(), reason="MPS device not available")
+class TestMPSDtype:
+    """Tests that MPS-incompatible dtypes (float64) are downcast to float32 in tensor specs."""
+
+    def test_mps_does_not_support_float64(self):
+        """Assert that MPS still doesn't support float64.
+
+        If this test fails, MPS has gained float64 support and the downcasts
+        can be removed (e.g., in _default_dtype_and_device)
+        """
+        with pytest.raises(TypeError, match="MPS framework doesn't support float64"):
+            torch.ones(2, dtype=torch.float64, device="mps")
+
+    def test_unbounded_to_mps_downcasts_float64(self):
+        """Unbounded.to('mps') downcasts float64 -> float32."""
+        spec_cpu = Unbounded(shape=(6,), device="cpu", dtype=torch.float64)
+        assert spec_cpu.dtype == torch.float64
+
+        with pytest.warns(UserWarning, match="MPS device does not support float64"):
+            spec_mps = spec_cpu.to("mps")
+        assert spec_mps.dtype == torch.float32
+        assert spec_mps.device.type == "mps"
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)

torchrl/data/tensor_specs.py

@@ -113,6 +113,15 @@ def _default_dtype_and_device(
         device = _make_ordinal_device(torch.device(device))
     elif not allow_none_device:
         device = torch.zeros(()).device
+
+    if device is not None and device.type == "mps" and dtype == torch.float64:
+        warnings.warn(
+            "MPS device does not support float64. Downcasting dtype from float64 to float32.",
+            UserWarning,
+            stacklevel=2,
+        )
+        dtype = torch.float32
+
     return dtype, device