diff --git a/linear_operator/operators/cat_linear_operator.py b/linear_operator/operators/cat_linear_operator.py index d44463c4..803cbe4c 100644 --- a/linear_operator/operators/cat_linear_operator.py +++ b/linear_operator/operators/cat_linear_operator.py @@ -397,7 +397,8 @@ def inv_quad_logdet( Tensor | None, # shape: (...) ]: # fmt: on res = super().inv_quad_logdet(inv_quad_rhs, logdet, reduce_inv_quad) - return tuple(r.to(self.device) for r in res) + # a term that is not requested can be None + return tuple(r.to(self.device) if r is not None else None for r in res) @property def device(self) -> torch.device | None: diff --git a/linear_operator/operators/chol_linear_operator.py b/linear_operator/operators/chol_linear_operator.py index 697e784f..7c2e59d8 100644 --- a/linear_operator/operators/chol_linear_operator.py +++ b/linear_operator/operators/chol_linear_operator.py @@ -109,6 +109,9 @@ def inv_quad( inv_quad_rhs: Tensor, # shape: (*batch, N, M) or (*batch, N) reduce_inv_quad: bool = True, ) -> Tensor: # shape: (*batch, M) or (*batch) + if self.dim() == 2 and inv_quad_rhs.dim() == 1: + # a vector right-hand side is a single column, as in the default implementation + inv_quad_rhs = inv_quad_rhs.unsqueeze(-1) if self.upper: R = self.root._transpose_nonbatch().solve(inv_quad_rhs) else: diff --git a/test/operators/test_cat_linear_operator.py b/test/operators/test_cat_linear_operator.py index e652dfef..a2cbf069 100644 --- a/test/operators/test_cat_linear_operator.py +++ b/test/operators/test_cat_linear_operator.py @@ -28,6 +28,14 @@ def create_linear_op(self): def evaluate_linear_op(self, linear_op): return self.psd_mat.detach().clone().requires_grad_() + def test_logdet_without_inv_quad(self): + linear_op = self.create_linear_op().detach() + evaluated = self.evaluate_linear_op(linear_op).detach() + self.assertAllClose(linear_op.logdet(), torch.logdet(evaluated)) + inv_quad, logdet = linear_op.inv_quad_logdet(logdet=True) + self.assertIsNone(inv_quad) + self.assertAllClose(logdet, torch.logdet(evaluated)) + class TestCatLinearOperatorColumn(LinearOperatorTestCase, unittest.TestCase): seed = 1 diff --git a/test/operators/test_chol_linear_operator.py b/test/operators/test_chol_linear_operator.py index 52a7a5d2..6a4938ca 100644 --- a/test/operators/test_chol_linear_operator.py +++ b/test/operators/test_chol_linear_operator.py @@ -4,7 +4,7 @@ import torch -from linear_operator.operators import CholLinearOperator, TriangularLinearOperator +from linear_operator.operators import CholLinearOperator, DenseLinearOperator, TriangularLinearOperator from linear_operator.test.linear_operator_test_case import LinearOperatorTestCase @@ -33,6 +33,23 @@ def evaluate_linear_op(self, linear_op): chol = linear_op.root.to_dense() return chol.matmul(chol.mT) + def test_inv_quad_vector(self): + linear_op = self.create_linear_op().detach() + if linear_op.batch_shape: + return # vector right-hand sides are only accepted for non-batch operators + evaluated = self.evaluate_linear_op(linear_op) + vec = torch.randn(linear_op.size(-1)) + expected = vec @ torch.linalg.solve(evaluated, vec) + self.assertAllClose(linear_op.inv_quad(vec), expected) + self.assertAllClose(linear_op.inv_quad(vec, reduce_inv_quad=False), expected.unsqueeze(-1)) + inv_quad, logdet = linear_op.inv_quad_logdet(vec, logdet=True) + self.assertAllClose(inv_quad, expected) + self.assertAllClose(logdet, torch.logdet(evaluated)) + # dense operators of this size are handled through a CholLinearOperator as well + inv_quad, logdet = DenseLinearOperator(evaluated).inv_quad_logdet(vec, logdet=True) + self.assertAllClose(inv_quad, expected) + self.assertAllClose(logdet, torch.logdet(evaluated)) + def test_inverse(self): linear_op = self.create_linear_op() linear_op_copy = linear_op.detach().clone()