Fix inv_quad_logdet with a vector rhs (Chol/Dense) and CatLinearOperator.logdet()

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:

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:

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

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()