Fix add_low_rank to only compute roots when cached roots exist

This fixes a numerical instability bug where add_low_rank would speculatively
compute root decompositions even when no cached roots existed. This caused
SVD failures on ill-conditioned matrices (e.g., when using LinearKernel).

Changes:
- Change early exit condition from 'not generate_roots and not has_roots'
  to 'not (generate_roots and has_roots)' - exit early if EITHER
  generate_roots=False OR no cached roots exist
- Pass generate_roots parameter through in RootLinearOperator.add_low_rank
  (was being silently ignored)
- Add regression test to verify root_decomposition is not called when
  no roots are cached
- Update test_cat_rows to first cache roots before expecting them to be
  cached after cat_rows (aligns test with new expected behavior)

Author: SebastianAment

linear_operator/operators/_linear_operator.py

@@ -1074,9 +1074,10 @@ def add_low_rank(
                 new_linear_op = to_linear_operator(new_linear_op.to_dense())
 
         # if the old LinearOperator does not have either a root decomposition or a root inverse decomposition
-        # don't create one
+        # don't create one. Also skip if the caller explicitly doesn't want roots generated.
+        # The root update is only beneficial when self already has cached roots that can be efficiently updated.
         has_roots = any(_is_in_cache_ignore_args(self, key) for key in ("root_decomposition", "root_inv_decomposition"))
-        if not generate_roots and not has_roots:
+        if not (generate_roots and has_roots):
             return new_linear_op
 
         # we are going to compute the following
@@ -1253,15 +1254,16 @@ def cat_rows(
         new_linear_op = CatLinearOperator(upper_row, lower_row, dim=-1, output_device=A.device)
 
         # if the old LinearOperator does not have either a root decomposition or a root inverse decomposition
-        # don't create one
+        # don't create one. Also skip if the caller explicitly doesn't want roots generated.
+        # The root update is only beneficial when self already has cached roots that can be efficiently updated.
         has_roots = any(
             _is_in_cache_ignore_args(self, key)
             for key in (
                 "root_decomposition",
                 "root_inv_decomposition",
             )
         )
-        if not generate_roots and not has_roots:
+        if not (generate_roots and has_roots):
             return new_linear_op
 
         # Get components for new root Z = [E 0; F G]

linear_operator/operators/root_linear_operator.py

@@ -95,7 +95,11 @@ def add_low_rank(
         generate_roots: bool | None = True,
         **root_decomp_kwargs,
     ) -> LinearOperator:  # shape: (*batch, N, N)
-        return super().add_low_rank(low_rank_mat, root_inv_decomp_method=root_inv_decomp_method)
+        return super().add_low_rank(
+            low_rank_mat,
+            root_inv_decomp_method=root_inv_decomp_method,
+            generate_roots=generate_roots,
+        )
 
     def root_decomposition(
         self: LinearOperator, method: str | None = None  # shape: (*batch, N, N)

linear_operator/test/linear_operator_test_case.py

@@ -733,8 +733,14 @@ def test_cat_rows(self):
             root_rhs = linear_operator.root_decomposition(new_lt).matmul(rhs)
             self.assertAllClose(root_rhs, concat_rhs, **self.tolerances["root_decomposition"])
 
-            # check that root inv is cached
-            root_inv = get_from_cache(new_lt, "root_inv_decomposition")
+            # Test root_inv caching: roots are only updated when cached roots already exist.
+            # First, ensure linear_op has cached roots before calling cat_rows.
+            _ = linear_op.root_decomposition()
+            _ = linear_op.root_inv_decomposition()
+            new_lt_with_roots = linear_op.cat_rows(new_rows, new_point)
+
+            # check that root inv is cached (since linear_op had cached roots)
+            root_inv = get_from_cache(new_lt_with_roots, "root_inv_decomposition")
             # check that the inverse root decomposition is close
             concat_solve = torch.linalg.solve(concatenated_lt, rhs.unsqueeze(-1)).squeeze(-1)
             root_inv_solve = root_inv.matmul(rhs)

test/operators/test_dense_linear_operator.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 
 import unittest
+from unittest.mock import patch
 
 import torch
 
@@ -30,6 +31,48 @@ def test_root_decomposition_exact(self):
             actual = linear_op.matmul(test_mat)
             self.assertLess(torch.norm(res - actual) / actual.norm(), 0.1)
 
+    def test_no_root_computation_when_no_cached_roots(self):
+        """
+        Regression test for add_low_rank speculative root computation bug.
+        Verify root_decomposition is NOT called when no roots are cached.
+
+        This catches a bug where add_low_rank would unnecessarily compute expensive
+        root decompositions even when the base LinearOperator had no cached roots.
+        This caused numerical instability (SVD failures) on ill-conditioned matrices.
+
+        The fix ensures root updates only happen when BOTH:
+        1. generate_roots=True (default)
+        2. The base operator already has cached roots
+        """
+        torch.manual_seed(42)
+
+        # Create a simple PSD matrix without any cached root decomposition
+        n = 5
+        A = torch.randn(n, n)
+        base_matrix = A @ A.T + 0.1 * torch.eye(n)
+        base_op = DenseLinearOperator(base_matrix)
+
+        # Create a low-rank term (like LinearKernel produces)
+        low_rank = torch.randn(n, 2)
+
+        # Patch root_decomposition to track if it's called
+        # Before the fix, add_low_rank would call root_decomposition even when none are cached
+        # After the fix, it should NOT call root_decomposition
+        with patch.object(
+            DenseLinearOperator, "root_decomposition", wraps=base_op.root_decomposition
+        ) as mock_root_decomp:
+            result = base_op.add_low_rank(low_rank)
+
+            # Verify root_decomposition was NOT called (the fix's behavior)
+            # Before the fix, this would fail because root_decomposition was called
+            # add_low_rank should NOT compute root_decomposition when no roots are cached
+            self.assertEqual(mock_root_decomp.call_count, 0)
+
+        # Verify the result is still correct (simple matrix addition)
+        expected = base_matrix + low_rank @ low_rank.T
+        # add_low_rank should return correct sum
+        self.assertTrue(torch.allclose(result.to_dense(), expected, atol=1e-5))
+
 
 class TestDenseLinearOperatorBatch(LinearOperatorTestCase, unittest.TestCase):
     seed = 0