From 84f62571dede306b2693f78c8768f9a409e1c564 Mon Sep 17 00:00:00 2001 From: Zhe Wang <53119724+flyingdoubleG@users.noreply.github.com> Date: Thu, 5 Nov 2026 06:45:30 +0800 Subject: [PATCH] Fix inv_quad_logdet with a vector rhs and CatLinearOperator.logdet CholLinearOperator.inv_quad_logdet accepts a vector right-hand side for a non-batch operator, but inv_quad then summed over dim -2 of the 1-D solve and raised an IndexError. Since small operators compute inv_quad_logdet through a CholLinearOperator, this also broke e.g. DenseLinearOperator(A).inv_quad_logdet(vec, logdet=True). CatLinearOperator.inv_quad_logdet called .to(device) on every term, including the None that is returned when no right-hand side is given, so CatLinearOperator.logdet() raised an AttributeError. --- .../operators/cat_linear_operator.py | 3 ++- .../operators/chol_linear_operator.py | 3 +++ test/operators/test_cat_linear_operator.py | 8 ++++++++ test/operators/test_chol_linear_operator.py | 19 ++++++++++++++++++- 4 files changed, 31 insertions(+), 2 deletions(-) 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()