From cbd31958caaeeffa06a2158f7d650d68a1c682b2 Mon Sep 17 00:00:00 2001 From: Mark Colley <39366189+M-Colley@users.noreply.github.com> Date: Tue, 20 Oct 2026 20:34:14 +0200 Subject: [PATCH 1/2] Preserve triangular matmul higher-order gradients through analytic VJP --- .../operators/triangular_linear_operator.py | 5 +++ .../test_triangular_linear_operator.py | 39 +++++++++++++++++++ 2 files changed, 44 insertions(+) diff --git a/linear_operator/operators/triangular_linear_operator.py b/linear_operator/operators/triangular_linear_operator.py index fb5828b3..a3c2e29f 100644 --- a/linear_operator/operators/triangular_linear_operator.py +++ b/linear_operator/operators/triangular_linear_operator.py @@ -64,6 +64,11 @@ def __add__( return self.__class__(self._tensor + other._tensor, upper=self.upper) return self._tensor + other + def _bilinear_derivative(self, left_vecs: Tensor, right_vecs: Tensor) -> tuple[Tensor | None, ...]: + # The triangular wrapper does not change the representation. Preserve + # the wrapped operator's analytic derivative and its autograd graph. + return self._tensor._bilinear_derivative(left_vecs, right_vecs) + def _cholesky( self: LinearOperator, upper: bool | None = False # shape: (*batch, N, N) ) -> LinearOperator: # shape: (*batch, N, N) diff --git a/test/operators/test_triangular_linear_operator.py b/test/operators/test_triangular_linear_operator.py index 3bac5649..d7c4fd3b 100644 --- a/test/operators/test_triangular_linear_operator.py +++ b/test/operators/test_triangular_linear_operator.py @@ -15,6 +15,45 @@ class TestTriangularLinearOperator(LinearOperatorTestCase, unittest.TestCase): should_call_cg = False should_call_lanczos = False + def test_matmul_gradgrad(self): + for upper in (False, True): + for vector in (False, True): + with self.subTest(upper=upper, vector=vector): + raw = torch.tensor( + [[1.2, 0.1], [0.3, 0.9]], dtype=torch.double, requires_grad=True + ) + rhs = torch.tensor([[0.7, 0.2], [-0.4, 0.6]], dtype=torch.double) + rhs = (rhs[:, 0] if vector else rhs).requires_grad_() + + def matmul(raw, rhs): + triangle = raw.triu() if upper else raw.tril() + return TriangularLinearOperator(triangle, upper=upper).matmul( + rhs + ) + + self.assertTrue(torch.autograd.gradcheck(matmul, (raw, rhs))) + self.assertTrue(torch.autograd.gradgradcheck(matmul, (raw, rhs))) + + def test_cholesky_matmul_second_derivative(self): + raw = torch.tensor( + [[1.2, 0.1], [0.3, 0.9]], dtype=torch.double, requires_grad=True + ) + rhs = torch.tensor([0.7, -0.4], dtype=torch.double) + gradients = [] + for wrapped in (False, True): + covariance = raw @ raw.mT + torch.eye(2, dtype=torch.double) + root = torch.linalg.cholesky(covariance) + sample = ( + TriangularLinearOperator(root).matmul(rhs) if wrapped else root @ rhs + ) + first = torch.autograd.grad(sample.square().sum(), raw, create_graph=True)[ + 0 + ] + second = torch.autograd.grad(first.sum(), raw)[0] + gradients.append((first, second)) + for actual, expected in zip(*gradients): + self.assertAllClose(actual, expected) + def _ensure_symmetric_grad(self, grad): # Hack! we don't actually want symmetric grads for this LinearOperator test case # We actually want triangular gradients From 757c6e1e8cd49b6b685135eade508e92dbc2501c Mon Sep 17 00:00:00 2001 From: Mark Colley <39366189+M-Colley@users.noreply.github.com> Date: Tue, 20 Oct 2026 20:41:39 +0200 Subject: [PATCH 2/2] Preserve higher-order triangular matmul derivatives via analytic VJP --- .../operators/triangular_linear_operator.py | 1 - .../test_triangular_linear_operator.py | 43 ++++++++++++------- 2 files changed, 28 insertions(+), 16 deletions(-) diff --git a/linear_operator/operators/triangular_linear_operator.py b/linear_operator/operators/triangular_linear_operator.py index a3c2e29f..1a875585 100644 --- a/linear_operator/operators/triangular_linear_operator.py +++ b/linear_operator/operators/triangular_linear_operator.py @@ -9,7 +9,6 @@ from linear_operator.operators._linear_operator import IndexType, LinearOperator from linear_operator.operators.batch_repeat_linear_operator import BatchRepeatLinearOperator from linear_operator.operators.dense_linear_operator import DenseLinearOperator - from linear_operator.utils.errors import NotPSDError from linear_operator.utils.memoize import cached diff --git a/test/operators/test_triangular_linear_operator.py b/test/operators/test_triangular_linear_operator.py index d7c4fd3b..766d5bc6 100644 --- a/test/operators/test_triangular_linear_operator.py +++ b/test/operators/test_triangular_linear_operator.py @@ -19,41 +19,54 @@ def test_matmul_gradgrad(self): for upper in (False, True): for vector in (False, True): with self.subTest(upper=upper, vector=vector): - raw = torch.tensor( - [[1.2, 0.1], [0.3, 0.9]], dtype=torch.double, requires_grad=True - ) + raw = torch.tensor([[1.2, 0.1], [0.3, 0.9]], dtype=torch.double, requires_grad=True) rhs = torch.tensor([[0.7, 0.2], [-0.4, 0.6]], dtype=torch.double) rhs = (rhs[:, 0] if vector else rhs).requires_grad_() def matmul(raw, rhs): triangle = raw.triu() if upper else raw.tril() - return TriangularLinearOperator(triangle, upper=upper).matmul( - rhs - ) + return TriangularLinearOperator(triangle, upper=upper).matmul(rhs).square() self.assertTrue(torch.autograd.gradcheck(matmul, (raw, rhs))) self.assertTrue(torch.autograd.gradgradcheck(matmul, (raw, rhs))) def test_cholesky_matmul_second_derivative(self): - raw = torch.tensor( - [[1.2, 0.1], [0.3, 0.9]], dtype=torch.double, requires_grad=True - ) + raw = torch.tensor([[1.2, 0.1], [0.3, 0.9]], dtype=torch.double, requires_grad=True) rhs = torch.tensor([0.7, -0.4], dtype=torch.double) gradients = [] for wrapped in (False, True): covariance = raw @ raw.mT + torch.eye(2, dtype=torch.double) root = torch.linalg.cholesky(covariance) - sample = ( - TriangularLinearOperator(root).matmul(rhs) if wrapped else root @ rhs - ) - first = torch.autograd.grad(sample.square().sum(), raw, create_graph=True)[ - 0 - ] + sample = TriangularLinearOperator(root).matmul(rhs) if wrapped else root @ rhs + first = torch.autograd.grad(sample.square().sum(), raw, create_graph=True)[0] second = torch.autograd.grad(first.sum(), raw)[0] gradients.append((first, second)) for actual, expected in zip(*gradients): self.assertAllClose(actual, expected) + def test_matrix_hessian_with_constant_rhs(self): + devices = [torch.device("cpu")] + if torch.cuda.is_available(): + devices.append(torch.device("cuda")) + for device in devices: + with self.subTest(device=device): + root = torch.tensor( + [[1.2, 0.0], [0.3, 0.9]], + dtype=torch.double, + device=device, + requires_grad=True, + ) + rhs = root.new_tensor([0.7, -0.4]) + wrapped = TriangularLinearOperator(root).matmul(rhs) + first = torch.autograd.grad(wrapped.square().sum(), root, create_graph=True)[0] + expected = torch.autograd.grad((root @ rhs).square().sum(), root, create_graph=True)[0] + self.assertAllClose(first, expected) + self.assertTrue(first.requires_grad) + self.assertAllClose( + torch.autograd.grad(first.sum(), root)[0], + torch.autograd.grad(expected.sum(), root)[0], + ) + def _ensure_symmetric_grad(self, grad): # Hack! we don't actually want symmetric grads for this LinearOperator test case # We actually want triangular gradients