From c0326ec20bb3c201b1038ba1f55c8ea24108b230 Mon Sep 17 00:00:00 2001 From: Mark Colley <39366189+M-Colley@users.noreply.github.com> Date: Wed, 21 Oct 2026 17:34:26 +0200 Subject: [PATCH] Preserve higher derivatives in fallback bilinear autograd --- linear_operator/operators/_linear_operator.py | 13 ++- .../test_bilinear_derivative_autograd.py | 103 ++++++++++++++++++ 2 files changed, 111 insertions(+), 5 deletions(-) create mode 100644 test/operators/test_bilinear_derivative_autograd.py diff --git a/linear_operator/operators/_linear_operator.py b/linear_operator/operators/_linear_operator.py index dc830496..3b8d1fab 100644 --- a/linear_operator/operators/_linear_operator.py +++ b/linear_operator/operators/_linear_operator.py @@ -14,11 +14,9 @@ import numpy as np import torch - from torch import Tensor import linear_operator - from linear_operator import settings, utils from linear_operator.functions._diagonalization import Diagonalization from linear_operator.functions._inv_quad import InvQuad @@ -362,12 +360,15 @@ def _bilinear_derivative(self, left_vecs: Tensor, right_vecs: Tensor) -> tuple[T """ from collections import deque - # Construct a detached version of each argument in the linear operator + create_graph = torch.is_grad_enabled() + # Independent views make these partial derivatives with respect to the + # representation, even when representations or vector inputs depend on + # one another. Retain their original graph for higher-order backward. args = [] for arg in self.representation(): # All arguments here are guaranteed to be tensors if arg.dtype.is_floating_point and arg.requires_grad: - args.append(arg.detach().requires_grad_(True)) + args.append(arg.view_as(arg) if create_graph else arg.detach().requires_grad_(True)) else: args.append(arg.detach()) @@ -380,7 +381,9 @@ def _bilinear_derivative(self, left_vecs: Tensor, right_vecs: Tensor) -> tuple[T lin_op = self.representation_tree()(*args) loss = (left_vecs * lin_op._matmul(right_vecs)).sum() args_with_grads = [arg for arg in args if arg.requires_grad] - actual_grads = deque(torch.autograd.grad(loss, args_with_grads, allow_unused=True)) + actual_grads = deque( + torch.autograd.grad(loss, args_with_grads, allow_unused=True, create_graph=create_graph) + ) # Now make sure that the object we return has one entry for every item in args grads = [] diff --git a/test/operators/test_bilinear_derivative_autograd.py b/test/operators/test_bilinear_derivative_autograd.py new file mode 100644 index 00000000..b26ab3bf --- /dev/null +++ b/test/operators/test_bilinear_derivative_autograd.py @@ -0,0 +1,103 @@ +#!/usr/bin/env python3 + +import unittest + +import torch + +from linear_operator.operators import ( + DiagLinearOperator, + KroneckerProductLinearOperator, + LinearOperator, + RootLinearOperator, +) + + +class _NonlinearOperator(LinearOperator): + def __init__(self, root, diagonal): + super().__init__(root, diagonal) + self.root, self.diagonal = root, diagonal + + def _matmul(self, rhs): + return (self.root @ self.root.mT + self.diagonal.exp().diag_embed()) @ rhs + + def _size(self): + return self.root.shape[:-1] + self.root.shape[-2:-1] + + def _transpose_nonbatch(self): + return self + + +class TestBilinearDerivativeAutograd(unittest.TestCase): + def _devices(self): + return ("cpu", "cuda") if torch.cuda.is_available() else ("cpu",) + + def test_fallback_partial_derivative_and_mixed_higher_gradients(self): + for device in self._devices(): + with self.subTest(device=device): + root = torch.tensor([[0.3], [-0.2]], device=device, dtype=torch.double, requires_grad=True) + diagonal = root[:, 0].square() - 0.7 + left = root.sin().expand(2, 3) + right = root.cos().expand(2, 3) + operator = _NonlinearOperator(root, diagonal) + actual = operator._bilinear_derivative(left, right) + product = left @ right.mT + expected = ((product + product.mT) @ root, product.diagonal() * diagonal.exp()) + for observed, reference in zip(actual, expected): + torch.testing.assert_close(observed, reference) + self.assertTrue(observed.requires_grad) + actual_higher = torch.autograd.grad(sum(g.square().sum() for g in actual), root, retain_graph=True)[0] + expected_higher = torch.autograd.grad(sum(g.square().sum() for g in expected), root)[0] + torch.testing.assert_close(actual_higher, expected_higher) + + def test_kronecker_composition_matches_dense_through_third_derivatives(self): + for device in self._devices(): + for batch in ((), (2,), (2, 1)): + with self.subTest(device=device, batch=batch): + root = torch.tensor([[0.3], [-0.2]], device=device, dtype=torch.double) + root = root.expand(*batch, 2, 1).clone().requires_grad_() + diagonal = root.new_tensor([-0.7, -0.4]).requires_grad_() + data_diagonal = root.new_tensor([0.8, 1.1, 0.9]).requires_grad_() + task_covariance = RootLinearOperator(root) + DiagLinearOperator(diagonal.exp()) + operator = KroneckerProductLinearOperator(DiagLinearOperator(data_diagonal), task_covariance) + task = root @ root.mT + diagonal.exp().diag_embed() + dense = (data_diagonal.diag().reshape(3, 1, 3, 1) * task.unsqueeze(-3).unsqueeze(-2)).reshape( + *batch, 6, 6 + ) + torch.testing.assert_close(operator.to_dense(), dense) + parameters = (root, diagonal, data_diagonal) + rows = [] + for result in (operator.to_dense(), dense): + first = torch.autograd.grad( + result.sin().sum(), parameters, create_graph=True, retain_graph=True + ) + self.assertTrue(all(g.requires_grad for g in first)) + second = torch.autograd.grad( + sum(g.square().sum() for g in first), parameters, create_graph=True, retain_graph=True + ) + third = torch.autograd.grad( + sum(g.square().sum() for g in second), parameters, retain_graph=True + ) + rows.append((first, second, third)) + for observed, reference in zip(*rows): + for a, b in zip(observed, reference): + torch.testing.assert_close(a, b) + + def test_fallback_gradcheck_gradgradcheck_and_ordinary_backward(self): + for device in self._devices(): + with self.subTest(device=device): + root = torch.tensor([[0.3], [-0.2]], device=device, dtype=torch.double, requires_grad=True) + diagonal = root.new_tensor([-0.7, -0.4], requires_grad=True) + rhs = root.new_tensor([[0.4, 0.7], [-0.2, 0.3]], requires_grad=True) + + def apply(root, diagonal, rhs): + return _NonlinearOperator(root, diagonal).matmul(rhs) + + self.assertTrue(torch.autograd.gradcheck(apply, (root, diagonal, rhs))) + self.assertTrue(torch.autograd.gradgradcheck(apply, (root, diagonal, rhs))) + loss = apply(root, diagonal, rhs).square().sum() + ordinary = torch.autograd.grad(loss, (root, diagonal, rhs)) + self.assertFalse(any(g.requires_grad for g in ordinary)) + + +if __name__ == "__main__": + unittest.main()