Keep the dtype and device in to() if only one of them is given

linear_operator/operators/cat_linear_operator.py

@@ -424,7 +424,7 @@ def to(
         """
         device, dtype = _to_helper(*args, **kwargs)
 
-        new_kwargs = {**self._kwargs, "output_device": device}
+        new_kwargs = {**self._kwargs, "output_device": self.output_device if device is None else device}
         res = self.__class__(*self._args, **new_kwargs)
 
         if dtype is not None:

linear_operator/operators/identity_linear_operator.py

@@ -279,7 +279,11 @@ def to(
         new_kwargs = {}
         for arg in self._args:
             if hasattr(arg, "to"):
-                if hasattr(arg, "dtype") and arg.dtype.is_floating_point == dtype.is_floating_point:
+                if (
+                    dtype is not None
+                    and hasattr(arg, "dtype")
+                    and arg.dtype.is_floating_point == dtype.is_floating_point
+                ):
                     new_args.append(arg.to(dtype=dtype, device=device))
                 else:
                     new_args.append(arg.to(device=device))
@@ -290,6 +294,6 @@ def to(
                 new_kwargs[name] = val.to(dtype=dtype, device=device)
             else:
                 new_kwargs[name] = val
-        new_kwargs["device"] = device
-        new_kwargs["dtype"] = dtype
+        new_kwargs["device"] = self.device if device is None else device
+        new_kwargs["dtype"] = self.dtype if dtype is None else dtype
         return self.__class__(*new_args, **new_kwargs)

linear_operator/operators/interpolated_linear_operator.py

@@ -476,7 +476,11 @@ def to(
         new_kwargs = {}
         for arg in self._args:
             if hasattr(arg, "to"):
-                if hasattr(arg, "dtype") and arg.dtype.is_floating_point == dtype.is_floating_point:
+                if (
+                    dtype is not None
+                    and hasattr(arg, "dtype")
+                    and arg.dtype.is_floating_point == dtype.is_floating_point
+                ):
                     new_args.append(arg.to(dtype=dtype, device=device))
                 else:
                     new_args.append(arg.to(device=device))

linear_operator/operators/masked_linear_operator.py

@@ -138,7 +138,11 @@ def to(
         new_kwargs = {}
         for arg in self._args:
             if hasattr(arg, "to"):
-                if hasattr(arg, "dtype") and arg.dtype.is_floating_point == dtype.is_floating_point:
+                if (
+                    dtype is not None
+                    and hasattr(arg, "dtype")
+                    and arg.dtype.is_floating_point == dtype.is_floating_point
+                ):
                     new_args.append(arg.to(dtype=dtype, device=device))
                 else:
                     new_args.append(arg.to(device=device))

linear_operator/test/linear_operator_test_case.py

@@ -1284,3 +1284,16 @@ def test_to_double(self):
         except RuntimeError:
             raise RuntimeError(f"Could not convert {type(linear_op)} to double.")
         self.assertEqual(linear_op.dtype, torch.float64)
+
+    def test_to_device_or_dtype_only(self):
+        # to() with only a device (only a dtype) should keep the dtype (device) of the linear_op
+        linear_op = self.create_linear_op()
+        evaluated = self.evaluate_linear_op(linear_op)
+        device = evaluated.device
+        for res in (linear_op.to(device), linear_op.to(device=device)):
+            self.assertEqual(res.dtype, linear_op.dtype)
+            self.assertEqual(res.device, device)
+            self.assertAllClose(res.to_dense(), evaluated)
+        res = res.to(torch.float64)
+        self.assertEqual(res.dtype, torch.float64)
+        self.assertEqual(res.device, device)

test/operators/test_interpolated_linear_operator.py

@@ -4,7 +4,7 @@
 
 import torch
 
-from linear_operator.operators import DenseLinearOperator, InterpolatedLinearOperator
+from linear_operator.operators import CatLinearOperator, DenseLinearOperator, InterpolatedLinearOperator
 from linear_operator.test.linear_operator_test_case import LinearOperatorTestCase, RectangularLinearOperatorTestCase
 
 
@@ -61,6 +61,20 @@ def test_to_double(self):
         self.assertFalse(linear_op.left_interp_indices.dtype.is_floating_point)
         self.assertFalse(linear_op.right_interp_indices.dtype.is_floating_point)
 
+    def test_to_device_or_dtype_only(self):
+        super().test_to_device_or_dtype_only()
+        # check that the index matrices are still integer after conversion
+        linear_op = self.create_linear_op().to(torch.device("cpu"))
+        self.assertFalse(linear_op.left_interp_indices.dtype.is_floating_point)
+        self.assertFalse(linear_op.right_interp_indices.dtype.is_floating_point)
+
+    def test_getitem_cat(self):
+        # Indexing a CatLinearOperator moves the selected operator to the output device
+        linear_op = self.create_linear_op()
+        evaluated = self.evaluate_linear_op(linear_op)
+        cat = CatLinearOperator(linear_op.unsqueeze(0), linear_op.unsqueeze(0), dim=0, output_device=linear_op.device)
+        self.assertAllClose(cat[1].to_dense(), evaluated)
+
 
 class TestInterpolatedLinearOperatorBatch(LinearOperatorTestCase, unittest.TestCase):
     seed = 0

test/operators/test_masked_linear_operator.py

@@ -38,6 +38,13 @@ def test_to_double(self):
         self.assertEqual(linear_op.col_mask.dtype, torch.bool)
         self.assertEqual(linear_op.row_mask.dtype, torch.bool)
 
+    def test_to_device_or_dtype_only(self):
+        super().test_to_device_or_dtype_only()
+        # check that the mask matrices are still boolean after conversion
+        linear_op = self.create_linear_op().to(torch.device("cpu"))
+        self.assertEqual(linear_op.col_mask.dtype, torch.bool)
+        self.assertEqual(linear_op.row_mask.dtype, torch.bool)
+
 
 class TestMaskedLinearOperatorBatch(LinearOperatorTestCase, unittest.TestCase):
     seed = 2023