diff --git a/linear_operator/operators/zero_linear_operator.py b/linear_operator/operators/zero_linear_operator.py index cae129ee..1787ba46 100644 --- a/linear_operator/operators/zero_linear_operator.py +++ b/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,7 +226,7 @@ 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() @@ -235,7 +234,7 @@ def transpose(self, dim1: int, dim2: int) -> LinearOperator: 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) diff --git a/test/operators/test_zero_linear_operator.py b/test/operators/test_zero_linear_operator.py index 33ec9658..2b8c3348 100644 --- a/test/operators/test_zero_linear_operator.py +++ b/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)