Fix higher-order sparse solve gradients and optional JAX test skipping (#85)

Fix higher-order autograd support for sparse_generic_solve by routing the backward transpose solve through sparse_generic_solve rather than directly calling the raw iterative solver.

This avoids differentiating through solver internals when create_graph=True, while preserving the implicit-gradient path used by the custom autograd function.

Also:
- preserve the solve output in the autograd context without detaching it;
- add higher-order regression tests for linear_cg, minres, and BiCGSTAB;
- add a non-symmetric BiCGSTAB test with an explicit transpose-solve wrapper;
- fix optional JAX test skipping by moving the JAX import behind pytest.importorskip;
- bump package/docs version to 0.2.3.

Author: theo-barfoot

docs/source/conf.py

@@ -15,8 +15,8 @@
 project = "torchsparsegradutils"
 copyright = "2026, CAI4CAI research group"
 author = "CAI4CAI research group"
-release = "0.2.2"
-version = "0.2.2"
+release = "0.2.3"
+version = "0.2.3"
 
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "torchsparsegradutils"
-version = "0.2.2"
+version = "0.2.3"
 description = "A collection of utility functions to work with PyTorch sparse tensors"
 readme = "README.md"
 requires-python = ">=3.10"

torchsparsegradutils/sparse_solve.py

@@ -438,6 +438,7 @@ class SparseGenericSolve(torch.autograd.Function):
     @staticmethod
     def forward(ctx, A, B, solve, transpose_solve, kwargs):
         grad_flag = A.requires_grad or B.requires_grad
+        ctx.solve = solve
         ctx.transpose_solve = transpose_solve
         ctx.kwargs = kwargs  # Store kwargs for backward pass
 
@@ -449,7 +450,7 @@ def forward(ctx, A, B, solve, transpose_solve, kwargs):
 
         x.requires_grad = grad_flag
 
-        ctx.save_for_backward(A, x.detach())
+        ctx.save_for_backward(A, x)
         return x
 
     @staticmethod
@@ -463,7 +464,13 @@ def backward(ctx, grad):  # type: ignore[override]
             grad = grad.unsqueeze(-1)
 
         # Backprop rule: gradB = A^{-T} grad
-        gradB = ctx.transpose_solve(A, grad, **ctx.kwargs)
+        gradB = sparse_generic_solve(
+            A,
+            grad,
+            solve=ctx.transpose_solve,
+            transpose_solve=ctx.solve,
+            **ctx.kwargs,
+        )
 
         # Ensure gradient dtype matches input dtype
         if gradB.dtype != A.dtype:

torchsparsegradutils/tests/test_jax_bindings.py

@@ -1,19 +1,15 @@
-import jax
 import numpy as np
 import pytest
 import torch
 from test_config import DEVICES
 
-import torchsparsegradutils as tsgu
 import torchsparsegradutils.jax as tsgujax
 
 # skip if JAX unavailable
-pytest.importorskip("jax")
+jax = pytest.importorskip("jax")
 if not tsgujax.have_jax:
     pytest.skip("JAX bindings unavailable, skipping jax tests", allow_module_level=True)
 
-import jax.numpy as jnp
-
 
 def _id_device(d):
     return str(d)

torchsparsegradutils/tests/test_sparse_solve.py

@@ -42,6 +42,132 @@ def solve_id(solve):
     return "default"
 
 
+def _make_differentiable_tridiag_spd(theta, layout):
+    """Create a small sparse SPD matrix with fixed sparsity and differentiable values."""
+    n = theta.numel()
+    device = theta.device
+
+    diag = theta.square() + 2.0
+    off = -0.1 * torch.sigmoid(theta[:-1])
+
+    rows = torch.cat(
+        [
+            torch.arange(n, device=device),
+            torch.arange(n - 1, device=device),
+            torch.arange(1, n, device=device),
+        ]
+    )
+    cols = torch.cat(
+        [
+            torch.arange(n, device=device),
+            torch.arange(1, n, device=device),
+            torch.arange(n - 1, device=device),
+        ]
+    )
+    values = torch.cat([diag, off, off])
+
+    A = torch.sparse_coo_tensor(torch.stack([rows, cols]), values, (n, n)).coalesce()
+    if layout == torch.sparse_csr:
+        A = A.to_sparse_csr()
+    return A, A.to_dense()
+
+
+def _make_differentiable_nonsymmetric_tridiag(theta, layout):
+    """Create a small fixed-sparsity non-symmetric diagonally dominant matrix."""
+    n = 6
+    device = theta.device
+    dtype = theta.dtype
+
+    if theta.numel() != 3 * n - 2:
+        raise ValueError(f"theta should have length {3 * n - 2}, got {theta.numel()}")
+
+    diag_theta = theta[:n]
+    upper_theta = theta[n : n + n - 1]
+    lower_theta = theta[n + n - 1 :]
+
+    diag = diag_theta.square() + 3.0
+    upper = 0.05 * torch.tanh(upper_theta)
+    lower = -0.08 * torch.sigmoid(lower_theta)
+
+    rows = torch.cat(
+        [
+            torch.arange(n, device=device),
+            torch.arange(n - 1, device=device),
+            torch.arange(1, n, device=device),
+        ]
+    )
+    cols = torch.cat(
+        [
+            torch.arange(n, device=device),
+            torch.arange(1, n, device=device),
+            torch.arange(n - 1, device=device),
+        ]
+    )
+    values = torch.cat([diag, upper, lower]).to(dtype=dtype)
+
+    A = torch.sparse_coo_tensor(torch.stack([rows, cols]), values, (n, n)).coalesce()
+    if layout == torch.sparse_csr:
+        A = A.to_sparse_csr()
+    return A, A.to_dense()
+
+
+def _bicgstab_transpose(A, B, **kwargs):
+    """Solve A.T X = B using BiCGSTAB."""
+    if A.layout == torch.sparse_csr:
+        # A.T currently triggers aten::as_strided for sparse CSR tensors in PyTorch,
+        # so use transpose(...).to_sparse_csr() for the CSR test case.
+        # A.transpose(0, 1) for csr will return csc, hence the to_sparse_csr() call.
+        return bicgstab(A.transpose(0, 1).to_sparse_csr(), B, **kwargs)
+    return bicgstab(A.T, B, **kwargs)
+
+
+def _bicgstab_higher_order_kwargs(value_dtype):
+    from torchsparsegradutils.utils.bicgstab import BICGSTABSettings
+
+    if value_dtype == torch.float32:
+        return {"settings": BICGSTABSettings(reltol=1e-6, abstol=1e-6, matvec_max=1000)}
+    return {"settings": BICGSTABSettings(reltol=1e-12, abstol=1e-12, matvec_max=1000)}
+
+
+def _higher_order_bicgstab_tolerances(value_dtype):
+    if value_dtype == torch.float32:
+        return {
+            "output": (1e-5, 1e-5),
+            "grad": (5e-5, 5e-5),
+            "hess": (5e-3, 5e-3),
+        }
+    return {
+        "output": (1e-8, 1e-8),
+        "grad": (1e-6, 1e-6),
+        "hess": (5e-5, 5e-5),
+    }
+
+
+def _settings_for_higher_order_solve(solve, value_dtype):
+    if solve is linear_cg:
+        from torchsparsegradutils.utils.linear_cg import LinearCGSettings
+
+        return {"settings": LinearCGSettings(cg_tolerance=1e-5, max_cg_iterations=1000)}
+    if solve is minres:
+        from torchsparsegradutils.utils.minres import MINRESSettings
+
+        tolerance = 1e-6 if value_dtype == torch.float32 else 1e-10
+        return {"settings": MINRESSettings(minres_tolerance=tolerance, max_cg_iterations=1000)}
+    return {}
+
+
+def _higher_order_spd_tolerances(value_dtype):
+    if value_dtype == torch.float32:
+        return {
+            "grad": (1e-3, 1e-3),
+            "hess": (5e-2, 5e-2),
+        }
+    return {
+        "grad": (1e-5, 1e-5),
+        "hess": (5e-4, 5e-4),
+    }
+
+
 # Define Fixtures
 
 
@@ -260,3 +386,99 @@ def test_kwargs_with_different_solvers_same_matrix():
     # All solutions should be close to each other
     assert torch.allclose(X_cg, X_minres, atol=atol, rtol=rtol)
     assert torch.allclose(X_bicgstab, X_minres, atol=atol, rtol=rtol)
+
+
+@pytest.mark.parametrize("base_solve", [linear_cg, minres], ids=[solve_id(linear_cg), solve_id(minres)])
+def test_sparse_generic_solve_higher_order_create_graph_no_out_error(layout, base_solve, device, value_dtype):
+    torch.manual_seed(0)
+
+    theta = torch.randn(8, dtype=value_dtype, device=device, requires_grad=True)
+    A, _ = _make_differentiable_tridiag_spd(theta, layout)
+    B = torch.randn(8, 2, dtype=value_dtype, device=device)
+
+    kwargs = _settings_for_higher_order_solve(base_solve, value_dtype)
+    loss = sparse_generic_solve(A, B, solve=base_solve, transpose_solve=base_solve, **kwargs).sum()
+
+    grad_theta = torch.autograd.grad(loss, theta, create_graph=True)[0]
+
+    assert grad_theta.requires_grad
+    assert torch.isfinite(grad_theta).all()
+
+    second = torch.autograd.grad(grad_theta.sum(), theta)[0]
+    assert torch.isfinite(second).all()
+
+
+@pytest.mark.parametrize("base_solve", [linear_cg, minres], ids=[solve_id(linear_cg), solve_id(minres)])
+def test_sparse_generic_solve_higher_order_matches_dense_reference(layout, base_solve, device, value_dtype):
+    torch.manual_seed(1)
+
+    theta_sparse = torch.randn(6, dtype=value_dtype, device=device, requires_grad=True)
+    theta_dense = theta_sparse.detach().clone().requires_grad_()
+
+    B = torch.randn(6, 2, dtype=value_dtype, device=device)
+
+    A_sparse, _ = _make_differentiable_tridiag_spd(theta_sparse, layout)
+    _, A_dense = _make_differentiable_tridiag_spd(theta_dense, torch.sparse_coo)
+
+    kwargs = _settings_for_higher_order_solve(base_solve, value_dtype)
+    tolerances = _higher_order_spd_tolerances(value_dtype)
+
+    out_sparse = sparse_generic_solve(A_sparse, B, solve=base_solve, transpose_solve=base_solve, **kwargs)
+    out_dense = torch.linalg.solve(A_dense, B)
+
+    loss_sparse = out_sparse.square().sum()
+    loss_dense = out_dense.square().sum()
+
+    grad_sparse = torch.autograd.grad(loss_sparse, theta_sparse, create_graph=True)[0]
+    grad_dense = torch.autograd.grad(loss_dense, theta_dense, create_graph=True)[0]
+
+    hess_vec_sparse = torch.autograd.grad(grad_sparse.sum(), theta_sparse)[0]
+    hess_vec_dense = torch.autograd.grad(grad_dense.sum(), theta_dense)[0]
+
+    grad_atol, grad_rtol = tolerances["grad"]
+    hess_atol, hess_rtol = tolerances["hess"]
+    assert torch.allclose(grad_sparse, grad_dense, atol=grad_atol, rtol=grad_rtol)
+    assert torch.allclose(hess_vec_sparse, hess_vec_dense, atol=hess_atol, rtol=hess_rtol)
+
+
+def test_sparse_generic_solve_higher_order_nonsymmetric_bicgstab_matches_dense_reference(layout, device, value_dtype):
+    torch.manual_seed(2)
+
+    n = 6
+    theta_sparse = torch.randn(3 * n - 2, dtype=value_dtype, device=device, requires_grad=True)
+    theta_dense = theta_sparse.detach().clone().requires_grad_()
+
+    B = torch.randn(n, 2, dtype=value_dtype, device=device)
+
+    A_sparse, _ = _make_differentiable_nonsymmetric_tridiag(theta_sparse, layout)
+    _, A_dense = _make_differentiable_nonsymmetric_tridiag(theta_dense, torch.sparse_coo)
+    assert not torch.allclose(A_dense, A_dense.T)
+
+    kwargs = _bicgstab_higher_order_kwargs(value_dtype)
+    tolerances = _higher_order_bicgstab_tolerances(value_dtype)
+
+    out_sparse = sparse_generic_solve(
+        A_sparse,
+        B,
+        solve=bicgstab,
+        transpose_solve=_bicgstab_transpose,
+        **kwargs,
+    )
+    out_dense = torch.linalg.solve(A_dense, B)
+    output_atol, output_rtol = tolerances["output"]
+    assert torch.allclose(A_dense @ out_sparse, B, atol=output_atol, rtol=output_rtol)
+
+    loss_sparse = out_sparse.square().sum()
+    loss_dense = out_dense.square().sum()
+
+    grad_sparse = torch.autograd.grad(loss_sparse, theta_sparse, create_graph=True)[0]
+    grad_dense = torch.autograd.grad(loss_dense, theta_dense, create_graph=True)[0]
+
+    hess_vec_sparse = torch.autograd.grad(grad_sparse.sum(), theta_sparse)[0]
+    hess_vec_dense = torch.autograd.grad(grad_dense.sum(), theta_dense)[0]
+
+    grad_atol, grad_rtol = tolerances["grad"]
+    hess_atol, hess_rtol = tolerances["hess"]
+    assert torch.allclose(out_sparse, out_dense, atol=output_atol, rtol=output_rtol)
+    assert torch.allclose(grad_sparse, grad_dense, atol=grad_atol, rtol=grad_rtol)
+    assert torch.allclose(hess_vec_sparse, hess_vec_dense, atol=hess_atol, rtol=hess_rtol)