From 63f6af5cbe2551f16e45f07c3e026fedf7aa2c3e Mon Sep 17 00:00:00 2001 From: Mark Colley <39366189+M-Colley@users.noreply.github.com> Date: Mon, 12 Oct 2026 08:51:46 +0100 Subject: [PATCH] Keep the dtype and device in to() if only one of them is given _to_helper returns None for the dtype (device) if only a device (dtype) is passed to to(), but several to() overrides did not handle that: - InterpolatedLinearOperator.to and MaskedLinearOperator.to evaluated dtype.is_floating_point and raised "AttributeError: 'NoneType' object has no attribute 'is_floating_point'" for e.g. op.to(device="cpu"). This also broke indexing a CatLinearOperator of such operators, which calls .to(output_device) on the selected operator. - IdentityLinearOperator.to set the dtype (device) to None, so op.to(device) evaluated in the default dtype instead of op.dtype, and op.to(dtype) moved a CUDA operator to the CPU. - CatLinearOperator.to set output_device to None if only a dtype is given. Only move the device if no dtype is given, and keep the current dtype, device or output device if it is not given. Add a test to LinearOperatorTestCase that calls to() with only a device and only a dtype, and a test that indexes a CatLinearOperator of InterpolatedLinearOperators. Co-Authored-By: Claude Opus 5.5 --- linear_operator/operators/cat_linear_operator.py | 2 +- .../operators/identity_linear_operator.py | 10 +++++++--- .../operators/interpolated_linear_operator.py | 6 +++++- .../operators/masked_linear_operator.py | 6 +++++- .../test/linear_operator_test_case.py | 13 +++++++++++++ .../test_interpolated_linear_operator.py | 16 +++++++++++++++- test/operators/test_masked_linear_operator.py | 7 +++++++ 7 files changed, 53 insertions(+), 7 deletions(-) diff --git a/linear_operator/operators/cat_linear_operator.py b/linear_operator/operators/cat_linear_operator.py index d44463c4..a146b287 100644 --- a/linear_operator/operators/cat_linear_operator.py +++ b/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: diff --git a/linear_operator/operators/identity_linear_operator.py b/linear_operator/operators/identity_linear_operator.py index 54bf413e..28087752 100644 --- a/linear_operator/operators/identity_linear_operator.py +++ b/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) diff --git a/linear_operator/operators/interpolated_linear_operator.py b/linear_operator/operators/interpolated_linear_operator.py index 5a60ecb1..398a5675 100644 --- a/linear_operator/operators/interpolated_linear_operator.py +++ b/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)) diff --git a/linear_operator/operators/masked_linear_operator.py b/linear_operator/operators/masked_linear_operator.py index 782a4313..ef93fc51 100644 --- a/linear_operator/operators/masked_linear_operator.py +++ b/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)) diff --git a/linear_operator/test/linear_operator_test_case.py b/linear_operator/test/linear_operator_test_case.py index d33755bd..05db44c3 100644 --- a/linear_operator/test/linear_operator_test_case.py +++ b/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) diff --git a/test/operators/test_interpolated_linear_operator.py b/test/operators/test_interpolated_linear_operator.py index 5857e77b..197f7f35 100644 --- a/test/operators/test_interpolated_linear_operator.py +++ b/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 diff --git a/test/operators/test_masked_linear_operator.py b/test/operators/test_masked_linear_operator.py index 4c6c33aa..04f7aae5 100644 --- a/test/operators/test_masked_linear_operator.py +++ b/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