Preserve zero operator device, dtype and batch metadata

linear_operator/operators/zero_linear_operator.py

@@ -6,7 +6,6 @@
 from torch import Tensor
 
 from linear_operator.operators._linear_operator import IndexType, LinearOperator
-
 from linear_operator.utils.getitem import _compute_getitem_size
 from linear_operator.utils.memoize import cached
 
@@ -51,11 +50,11 @@ def _expand_batch(
 
     def _get_indices(self, row_index: IndexType, col_index: IndexType, *batch_indices: IndexType) -> torch.Tensor:
         new_size = _compute_getitem_size(self, batch_indices + (row_index, col_index))
-        return torch.zeros(*new_size)
+        return torch.zeros(new_size, dtype=self.dtype, device=self.device)
 
     def _getitem(self, row_index: IndexType, col_index: IndexType, *batch_indices: IndexType) -> LinearOperator:
         new_size = _compute_getitem_size(self, batch_indices + (row_index, col_index))
-        return ZeroLinearOperator(*new_size)
+        return ZeroLinearOperator(*new_size, dtype=self.dtype, device=self.device)
 
     def _matmul(
         self: LinearOperator,  # shape: (*batch, M, N)
@@ -191,7 +190,7 @@ def inv_quad_logdet(
     def logdet(
         self: LinearOperator,  # shape: (*batch, M, N)
     ) -> Tensor:  # shape: (*batch)
-        return torch.log(torch.tensor(0.0))
+        return torch.full(self.batch_shape, -torch.inf, dtype=self.dtype, device=self.device)
 
     def matmul(
         self: LinearOperator,  # shape: (*batch, M, N)
@@ -227,15 +226,15 @@ def solve(
     def to_dense(
         self: LinearOperator,  # shape: (*batch, M, N)
     ) -> Tensor:  # shape: (*batch, M, N)
-        return torch.zeros(*self.sizes)
+        return torch.zeros(self.sizes, dtype=self.dtype, device=self.device)
 
     def transpose(self, dim1: int, dim2: int) -> LinearOperator:
         sizes = self.sizes.copy()
         tmp = sizes[dim1]
         sizes[dim1] = sizes[dim2]
         sizes[dim2] = tmp
 
-        return ZeroLinearOperator(*sizes)
+        return ZeroLinearOperator(*sizes, dtype=self.dtype, device=self.device)
 
     def __add__(
         self: LinearOperator,  # shape: (..., #M, #N)

test/operators/test_zero_linear_operator.py

@@ -9,6 +9,29 @@
 
 
 class TestZeroLinearOperator(unittest.TestCase):
+    def test_dtype_device_preservation(self):
+        devices = [torch.device("cpu")]
+        if torch.cuda.is_available():
+            devices.append(torch.device("cuda"))
+        for device in devices:
+            for dtype in (torch.float32, torch.float64):
+                with self.subTest(device=device, dtype=dtype):
+                    zero = ZeroLinearOperator(2, 4, 4, dtype=dtype, device=device)
+                    dense = torch.zeros(2, 4, 4, dtype=dtype, device=device)
+                    indices = torch.tensor([0, 1], device=device)
+                    outputs = (
+                        (zero.to_dense(), dense),
+                        (zero[0].to_dense(), dense[0]),
+                        (zero[:, 1, :], dense[:, 1, :]),
+                        (zero[:, indices, indices], dense[:, indices, indices]),
+                        (zero[0, 1, 2], dense[0, 1, 2]),
+                        (zero.transpose(-1, -2).to_dense(), dense.mT),
+                        (zero.transpose(0, 1).to_dense(), dense.transpose(0, 1)),
+                        (zero.logdet(), torch.full((2,), -torch.inf, dtype=dtype, device=device)),
+                    )
+                    for actual, expected in outputs:
+                        torch.testing.assert_close(actual, expected)
+
     def test_to_dense(self):
         lv = ZeroLinearOperator(5, 4, 3)
         actual = torch.zeros(5, 4, 3)