torchsparsegradutils 0.2.4__tar.gz → 0.2.6__tar.gz

{torchsparsegradutils-0.2.4/torchsparsegradutils.egg-info → torchsparsegradutils-0.2.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torchsparsegradutils
-Version: 0.2.4
+Version: 0.2.6
 Summary: A collection of utility functions to work with PyTorch sparse tensors
 Author-email: CAI4CAI research group <contact@cai4cai.uk>
 License-Expression: Apache-2.0
@@ -239,15 +239,18 @@ batch_size, event_size = 2, 1000
 loc = torch.zeros(batch_size, event_size)
 
 # Example 1: LDL^T parameterization (numerically stable for precision matrices)
-# Create sparse lower triangular matrix (unit triangular, no diagonal)
+# Create sparse lower triangular matrix (strictly lower triangular, no diagonal)
+
 scale_tril = rand_sparse_tri(
     (batch_size, event_size, event_size),
-    nnz=5000,  # 5000 non-zeros for 1M parameters (0.5% sparsity)
+    nnz=5000,  # 5000 non-zeros for 1000x1000 dense matrix (99.5% sparsity)
     layout=torch.sparse_csr,
     upper=False,
-    unit_triangular=True  # Unit triangular for LDL^T
+    strict=True  # Strict triangular (exclude diagonal)
 )
 
+scale_tril.requires_grad_(True)
+
 # Diagonal component for LDL^T parameterization
 diagonal = torch.ones(batch_size, event_size) * 0.5
 
@@ -255,7 +258,7 @@ diagonal = torch.ones(batch_size, event_size) * 0.5
 dist_ldlt = SparseMultivariateNormal(
     loc=loc,
     diagonal=diagonal,
-    scale_tril=scale_tril  # Unit lower triangular
+    scale_tril=scale_tril  # Strictly lower triangular
 )
 
 # Example 2: LL^T parameterization (standard Cholesky)
@@ -264,7 +267,7 @@ scale_tril_chol = rand_sparse_tri(
     nnz=5000,
     layout=torch.sparse_csr,
     upper=False,
-    unit_triangular=False  # Include diagonal for LL^T
+    strict=False  # Include diagonal
 )
 
 # Create distribution with LL^T parameterization
@@ -279,7 +282,7 @@ precision_tril = rand_sparse_tri(
     nnz=5000,
     layout=torch.sparse_csr,
     upper=False,
-    unit_triangular=True
+    strict=True
 )
 
 precision_diagonal = torch.ones(batch_size, event_size) * 2.0
@@ -287,7 +290,7 @@ precision_diagonal = torch.ones(batch_size, event_size) * 2.0
 dist_precision = SparseMultivariateNormal(
     loc=loc,
     diagonal=precision_diagonal,
-    precision_tril=precision_tril  # Unit triangular precision factor
+    precision_tril=precision_tril  # Strictly triangular precision factor
 )
 
 # Sample with gradient support

{torchsparsegradutils-0.2.4 → torchsparsegradutils-0.2.6}/README.md

@@ -192,15 +192,18 @@ batch_size, event_size = 2, 1000
 loc = torch.zeros(batch_size, event_size)
 
 # Example 1: LDL^T parameterization (numerically stable for precision matrices)
-# Create sparse lower triangular matrix (unit triangular, no diagonal)
+# Create sparse lower triangular matrix (strictly lower triangular, no diagonal)
+
 scale_tril = rand_sparse_tri(
     (batch_size, event_size, event_size),
-    nnz=5000,  # 5000 non-zeros for 1M parameters (0.5% sparsity)
+    nnz=5000,  # 5000 non-zeros for 1000x1000 dense matrix (99.5% sparsity)
     layout=torch.sparse_csr,
     upper=False,
-    unit_triangular=True  # Unit triangular for LDL^T
+    strict=True  # Strict triangular (exclude diagonal)
 )
 
+scale_tril.requires_grad_(True)
+
 # Diagonal component for LDL^T parameterization
 diagonal = torch.ones(batch_size, event_size) * 0.5
 
@@ -208,7 +211,7 @@ diagonal = torch.ones(batch_size, event_size) * 0.5
 dist_ldlt = SparseMultivariateNormal(
     loc=loc,
     diagonal=diagonal,
-    scale_tril=scale_tril  # Unit lower triangular
+    scale_tril=scale_tril  # Strictly lower triangular
 )
 
 # Example 2: LL^T parameterization (standard Cholesky)
@@ -217,7 +220,7 @@ scale_tril_chol = rand_sparse_tri(
     nnz=5000,
     layout=torch.sparse_csr,
     upper=False,
-    unit_triangular=False  # Include diagonal for LL^T
+    strict=False  # Include diagonal
 )
 
 # Create distribution with LL^T parameterization
@@ -232,7 +235,7 @@ precision_tril = rand_sparse_tri(
     nnz=5000,
     layout=torch.sparse_csr,
     upper=False,
-    unit_triangular=True
+    strict=True
 )
 
 precision_diagonal = torch.ones(batch_size, event_size) * 2.0
@@ -240,7 +243,7 @@ precision_diagonal = torch.ones(batch_size, event_size) * 2.0
 dist_precision = SparseMultivariateNormal(
     loc=loc,
     diagonal=precision_diagonal,
-    precision_tril=precision_tril  # Unit triangular precision factor
+    precision_tril=precision_tril  # Strictly triangular precision factor
 )
 
 # Sample with gradient support

{torchsparsegradutils-0.2.4 → torchsparsegradutils-0.2.6}/docs/source/conf.py

@@ -15,8 +15,8 @@ sys.path.insert(0, os.path.abspath("../../"))
 project = "torchsparsegradutils"
 copyright = "2026, CAI4CAI research group"
 author = "CAI4CAI research group"
-release = "0.2.4"
-version = "0.2.4"
+release = "0.2.6"
+version = "0.2.6"
 
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration

{torchsparsegradutils-0.2.4 → torchsparsegradutils-0.2.6}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "torchsparsegradutils"
-version = "0.2.4"
+version = "0.2.6"
 description = "A collection of utility functions to work with PyTorch sparse tensors"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -105,6 +105,12 @@ testpaths = [
 pythonpath = [
     "torchsparsegradutils/tests",
 ]
+# Fail tests when deprecated upstream APIs are exercised outside intentional compatibility paths.
+filterwarnings = [
+    'error:torch\.norm is deprecated.*:UserWarning',
+    'error:torch\.triangular_solve is deprecated.*:UserWarning',
+    'error:.*torch\.jit\.script.*is deprecated.*:DeprecationWarning',
+]
 doctest_optionflags = [
     "NORMALIZE_WHITESPACE",
     "IGNORE_EXCEPTION_DETAIL",

torchsparsegradutils-0.2.6/torchsparsegradutils/_compat.py

@@ -0,0 +1,48 @@
+"""Compatibility helpers for PyTorch APIs without feature-equivalent replacements."""
+
+import warnings
+
+import torch
+
+
+def linalg_solve_triangular_compat(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    *,
+    upper: bool,
+    unitriangular: bool = False,
+    transpose: bool = False,
+) -> torch.Tensor:
+    """Solve a triangular system with the appropriate dense or sparse backend.
+
+    Dense coefficient matrices use ``torch.linalg.solve_triangular``. Sparse
+    matrices continue to use the isolated legacy call because the replacement
+    API does not support the sparse layouts required by this package. See
+    https://github.com/pytorch/pytorch/issues/87358 for upstream sparse feature
+    parity tracking.
+    """
+    if A.layout == torch.strided:
+        if transpose:
+            A = A.transpose(-2, -1)
+            upper = not upper
+
+        return torch.linalg.solve_triangular(
+            A,
+            B,
+            upper=upper,
+            unitriangular=unitriangular,
+        )
+
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore",
+            message=r"torch\.triangular_solve is deprecated.*",
+            category=UserWarning,
+        )
+        return torch.triangular_solve(
+            B,
+            A,
+            upper=upper,
+            transpose=transpose,
+            unitriangular=unitriangular,
+        ).solution

torchsparsegradutils-0.2.6/torchsparsegradutils/benchmarks/linear_cg_convergence.py

@@ -0,0 +1,224 @@
+#!/usr/bin/env python3
+"""Compare linear-CG convergence before and after the diagnostics fix.
+
+This benchmark focuses on numerical behavior rather than throughput. It loads
+the historical implementation from Git so both solvers run in one Python
+environment on identical deterministic inputs.
+
+Example
+-------
+python -m torchsparsegradutils.benchmarks.linear_cg_convergence \
+    --baseline-ref ea7b8f0 --repeats 100
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import statistics
+import subprocess
+import time
+import types
+import warnings
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable
+
+import torch
+
+from torchsparsegradutils.utils import linear_cg
+
+REPOSITORY_ROOT = Path(__file__).resolve().parents[2]
+LINEAR_CG_PATH = "torchsparsegradutils/utils/linear_cg.py"
+
+
+@dataclass(frozen=True)
+class Problem:
+    name: str
+    matrix: torch.Tensor
+    rhs: torch.Tensor
+    tolerance: float
+    max_iter: int
+    initial_guess: torch.Tensor | None = None
+
+
+class CountedMatmul:
+    def __init__(self, matrix: torch.Tensor):
+        self.matrix = matrix
+        self.calls = 0
+
+    def __call__(self, value: torch.Tensor) -> torch.Tensor:
+        self.calls += 1
+        return self.matrix @ value
+
+
+def load_historical_linear_cg(revision: str) -> Callable[..., torch.Tensor]:
+    """Load ``linear_cg`` from a Git revision without changing the worktree."""
+    command = ["git", "-C", str(REPOSITORY_ROOT), "show", f"{revision}:{LINEAR_CG_PATH}"]
+    try:
+        completed = subprocess.run(command, check=True, capture_output=True, text=True)
+    except FileNotFoundError as error:
+        raise RuntimeError("Cannot load the historical solver because the Git executable was not found") from error
+    except subprocess.CalledProcessError as error:
+        detail = error.stderr.strip() or "Git could not resolve the requested revision and file"
+        raise RuntimeError(
+            f"Cannot load the historical solver from revision {revision!r}. "
+            f"Run this benchmark from a Git checkout containing that revision. Git reported: {detail}"
+        ) from error
+    module = types.ModuleType("historical_linear_cg")
+    exec(compile(completed.stdout, f"{revision}:{LINEAR_CG_PATH}", "exec"), module.__dict__)
+    return module.linear_cg
+
+
+def make_problems(device: torch.device) -> list[Problem]:
+    dtype = torch.float64
+
+    size = 36
+    diagonal = torch.full((size,), 2.0, dtype=dtype, device=device)
+    off_diagonal = torch.full((size - 1,), -0.25, dtype=dtype, device=device)
+    tight_matrix = torch.diag(diagonal) + torch.diag(off_diagonal, diagonal=1) + torch.diag(off_diagonal, diagonal=-1)
+    tight_rhs = torch.zeros(size, dtype=dtype, device=device)
+    tight_rhs[size // 2] = 1
+
+    size = 40
+    multiple_matrix = torch.diag(torch.linspace(1.0, 100.0, size, dtype=dtype, device=device))
+    multiple_rhs = torch.zeros((size, 101), dtype=dtype, device=device)
+    multiple_rhs[:, -1] = 1
+
+    zero_matrix = torch.diag(torch.tensor([1.0, 2.0, 3.0], dtype=dtype, device=device))
+    zero_rhs = torch.zeros(3, dtype=dtype, device=device)
+
+    return [
+        Problem("tight_tolerance", tight_matrix, tight_rhs, tolerance=1e-12, max_iter=200),
+        Problem("multiple_rhs", multiple_matrix, multiple_rhs, tolerance=1e-4, max_iter=size),
+        Problem(
+            "zero_rhs_nonzero_guess",
+            zero_matrix,
+            zero_rhs,
+            tolerance=1e-5,
+            max_iter=20,
+            initial_guess=torch.ones_like(zero_rhs),
+        ),
+    ]
+
+
+def synchronize(device: torch.device) -> None:
+    if device.type == "cuda":
+        torch.cuda.synchronize(device)
+
+
+def relative_residual_per_rhs(matrix: torch.Tensor, solution: torch.Tensor, rhs: torch.Tensor) -> torch.Tensor:
+    if rhs.ndim == 1:
+        rhs = rhs.unsqueeze(-1)
+        solution = solution.unsqueeze(-1)
+    residual_norm = torch.linalg.vector_norm(rhs - matrix @ solution, dim=-2)
+    rhs_norm = torch.linalg.vector_norm(rhs, dim=-2)
+    return residual_norm / rhs_norm.masked_fill(rhs_norm.eq(0), 1)
+
+
+def run_once(
+    solver: Callable[..., torch.Tensor],
+    problem: Problem,
+    *,
+    fixed: bool,
+) -> tuple[torch.Tensor, int, int, str]:
+    matmul = CountedMatmul(problem.matrix)
+    arguments = {
+        "tolerance": problem.tolerance,
+        "max_iter": problem.max_iter,
+        "initial_guess": problem.initial_guess,
+    }
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", UserWarning)
+        if fixed:
+            solution, info = solver(matmul, problem.rhs, return_info=True, **arguments)
+            return solution, info.iterations, matmul.calls, info.reason
+        solution = solver(matmul, problem.rhs, **arguments)
+    # The historical solver performs one initial matvec followed by one per
+    # iteration and does not recompute the true residual before returning.
+    return solution, max(0, matmul.calls - 1), matmul.calls, "not_reported"
+
+
+def measure(
+    solver: Callable[..., torch.Tensor],
+    problem: Problem,
+    *,
+    fixed: bool,
+    warmup: int,
+    repeats: int,
+) -> dict[str, object]:
+    for _ in range(warmup):
+        run_once(solver, problem, fixed=fixed)
+    synchronize(problem.rhs.device)
+
+    durations = []
+    for _ in range(repeats):
+        start = time.perf_counter()
+        run_once(solver, problem, fixed=fixed)
+        synchronize(problem.rhs.device)
+        durations.append(time.perf_counter() - start)
+
+    solution, iterations, matvecs, reason = run_once(solver, problem, fixed=fixed)
+    true_residual = relative_residual_per_rhs(problem.matrix, solution, problem.rhs)
+    reference = torch.linalg.solve(problem.matrix, problem.rhs)
+    error = torch.linalg.vector_norm(solution - reference) / torch.linalg.vector_norm(reference).clamp_min(1)
+
+    return {
+        "iterations": iterations,
+        "matvecs": matvecs,
+        "reason": reason,
+        "true_relative_residual_max": float(true_residual.max()),
+        "unconverged_rhs": int((true_residual > problem.tolerance).sum()),
+        "relative_solution_error": float(error),
+        "median_time_ms": statistics.median(durations) * 1e3,
+    }
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--baseline-ref", default="ea7b8f0", help="Git revision containing the historical solver")
+    parser.add_argument("--device", default="cpu", help="PyTorch device (default: cpu)")
+    parser.add_argument("--warmup", type=int, default=5)
+    parser.add_argument("--repeats", type=int, default=50)
+    arguments = parser.parse_args()
+    if arguments.warmup < 0 or arguments.repeats < 1:
+        parser.error("--warmup must be nonnegative and --repeats must be positive")
+    return arguments
+
+
+def main() -> None:
+    arguments = parse_args()
+    device = torch.device(arguments.device)
+    historical_linear_cg = load_historical_linear_cg(arguments.baseline_ref)
+    results = {}
+    for problem in make_problems(device):
+        results[problem.name] = {
+            "baseline": measure(
+                historical_linear_cg,
+                problem,
+                fixed=False,
+                warmup=arguments.warmup,
+                repeats=arguments.repeats,
+            ),
+            "fixed": measure(
+                linear_cg,
+                problem,
+                fixed=True,
+                warmup=arguments.warmup,
+                repeats=arguments.repeats,
+            ),
+        }
+
+    payload = {
+        "baseline_ref": arguments.baseline_ref,
+        "device": str(device),
+        "torch_version": torch.__version__,
+        "warmup": arguments.warmup,
+        "repeats": arguments.repeats,
+        "results": results,
+    }
+    print(json.dumps(payload, indent=2, sort_keys=True))
+
+
+if __name__ == "__main__":
+    main()

{torchsparsegradutils-0.2.4 → torchsparsegradutils-0.2.6}/torchsparsegradutils/benchmarks/sparse_generic_solve_rand.py

@@ -225,8 +225,8 @@ def run_sparse_generic_solve_benchmark():
                                 else:
                                     residual = A_sparse @ x - B
 
-                                resnorm = torch.norm(residual).cpu().item()
-                                B_norm = torch.norm(B).cpu().item()
+                                resnorm = torch.linalg.vector_norm(residual).cpu().item()
+                                B_norm = torch.linalg.vector_norm(B).cpu().item()
                                 relative_resnorm = resnorm / B_norm if B_norm > 0 else 0.0
 
                             # Print result with residual norm

{torchsparsegradutils-0.2.4 → torchsparsegradutils-0.2.6}/torchsparsegradutils/benchmarks/sparse_generic_solve_suite.py

@@ -201,10 +201,9 @@ def run_sparse_solve_benchmark():
                         with torch.no_grad():
                             x = alg_fn(A_sparse, B)
                             residual = A_sparse @ x - B
-                            resnorm = torch.norm(residual).cpu().item()
-                            relative_resnorm = (
-                                resnorm / torch.norm(B).cpu().item() if torch.norm(B).cpu().item() > 0 else 0.0
-                            )
+                            resnorm = torch.linalg.vector_norm(residual).cpu().item()
+                            B_norm = torch.linalg.vector_norm(B).cpu().item()
+                            relative_resnorm = resnorm / B_norm if B_norm > 0 else 0.0
 
                         # Print result with residual norm
                         print_result_row(

{torchsparsegradutils-0.2.4 → torchsparsegradutils-0.2.6}/torchsparsegradutils/benchmarks/sparse_triangular_solve_rand.py

@@ -27,6 +27,7 @@ from cupyx.scipy.sparse.linalg._solve import spsolve_triangular
 from tqdm import tqdm
 
 from torchsparsegradutils import sparse_triangular_solve
+from torchsparsegradutils._compat import linalg_solve_triangular_compat
 from torchsparsegradutils.cupy.cupy_sparse_solve import sparse_solve_c4t
 from torchsparsegradutils.utils import rand_sparse, rand_sparse_tri
 
@@ -58,16 +59,19 @@ LAYOUTS = [torch.sparse_coo, torch.sparse_csr]
 
 ALGORITHMS = [
     (
-        "dense.triangular_solve",
-        lambda A, B: torch.triangular_solve(
-            B, A.to_dense(), upper=UPPER, unitriangular=UNITRIANGULAR, transpose=TRANSPOSE
-        ).solution,
+        "dense.linalg.solve_triangular",
+        lambda A, B: torch.linalg.solve_triangular(
+            A.to_dense().transpose(-2, -1) if TRANSPOSE else A.to_dense(),
+            B,
+            upper=not UPPER if TRANSPOSE else UPPER,
+            unitriangular=UNITRIANGULAR,
+        ),
     ),
     (
-        "torch_triangular_solve",
-        lambda A, B: torch.triangular_solve(
-            B, A, upper=UPPER, unitriangular=UNITRIANGULAR, transpose=TRANSPOSE
-        ).solution,
+        "linalg_solve_triangular_compat",
+        lambda A, B: linalg_solve_triangular_compat(
+            A, B, upper=UPPER, unitriangular=UNITRIANGULAR, transpose=TRANSPOSE
+        ),
     ),
     (
         "sparse_triangular_solve",
@@ -153,8 +157,8 @@ def run_sparse_triangular_solve_benchmark():
                             with torch.no_grad():
                                 x = alg_fn(A_sparse, B)
                                 residual = A_sparse @ x - B
-                                resnorm = torch.norm(residual).cpu().item()
-                                relative_resnorm = resnorm / torch.norm(B).cpu().item()
+                                resnorm = torch.linalg.vector_norm(residual).cpu().item()
+                                relative_resnorm = resnorm / torch.linalg.vector_norm(B).cpu().item()
 
                             # # Calculate residual norm for solution accuracy
                             # with torch.no_grad():
@@ -165,8 +169,8 @@ def run_sparse_triangular_solve_benchmark():
                             #     else:
                             #         Ax = A_sparse @ x
                             #     residual = Ax - B
-                            #     resnorm = torch.norm(residual).cpu().item()
-                            #     relative_resnorm = resnorm / torch.norm(B).cpu().item()
+                            #     resnorm = torch.linalg.vector_norm(residual).cpu().item()
+                            #     relative_resnorm = resnorm / torch.linalg.vector_norm(B).cpu().item()
 
                             # Print result with residual norm
                             print_result_row(

{torchsparsegradutils-0.2.4 → torchsparsegradutils-0.2.6}/torchsparsegradutils/benchmarks/sparse_triangular_solve_suitesparse.py

@@ -34,6 +34,7 @@ from cupyx.scipy.sparse.linalg._solve import spsolve_triangular
 from tqdm import tqdm
 
 from torchsparsegradutils import sparse_triangular_solve
+from torchsparsegradutils._compat import linalg_solve_triangular_compat
 from torchsparsegradutils.cupy.cupy_sparse_solve import sparse_solve_c4t
 
 # from jax.lax.linalg import triangular_solve  # NOTE: jax doesn't have a sparse triangular solve
@@ -57,16 +58,19 @@ LAYOUTS = [torch.sparse_coo, torch.sparse_csr]
 
 ALGORITHMS = [
     (
-        "dense.triangular_solve",
-        lambda A, B: torch.triangular_solve(
-            B, A.to_dense(), upper=UPPER, unitriangular=UNITRIANGULAR, transpose=TRANSPOSE
-        ).solution,
+        "dense.linalg.solve_triangular",
+        lambda A, B: torch.linalg.solve_triangular(
+            A.to_dense().transpose(-2, -1) if TRANSPOSE else A.to_dense(),
+            B,
+            upper=not UPPER if TRANSPOSE else UPPER,
+            unitriangular=UNITRIANGULAR,
+        ),
     ),
     (
-        "torch_triangular_solve",
-        lambda A, B: torch.triangular_solve(
-            B, A, upper=UPPER, unitriangular=UNITRIANGULAR, transpose=TRANSPOSE
-        ).solution,
+        "linalg_solve_triangular_compat",
+        lambda A, B: linalg_solve_triangular_compat(
+            A, B, upper=UPPER, unitriangular=UNITRIANGULAR, transpose=TRANSPOSE
+        ),
     ),
     (
         "sparse_triangular_solve",
@@ -185,10 +189,9 @@ def run_triangular_solve_benchmark():
                             # Use the triangular matrix A_sparse (not A_full) for residual calculation
                             # This ensures the residual is computed correctly for the actual triangular system solved
                             residual = A_sparse @ x - B
-                            resnorm = torch.norm(residual).cpu().item()
-                            relative_resnorm = (
-                                resnorm / torch.norm(B).cpu().item() if torch.norm(B).cpu().item() > 0 else 0.0
-                            )
+                            resnorm = torch.linalg.vector_norm(residual).cpu().item()
+                            B_norm = torch.linalg.vector_norm(B).cpu().item()
+                            relative_resnorm = resnorm / B_norm if B_norm > 0 else 0.0
 
                         # Print result
                         print_result_row(

{torchsparsegradutils-0.2.4 → torchsparsegradutils-0.2.6}/torchsparsegradutils/sparse_solve.py

@@ -3,6 +3,7 @@ from typing import Callable, Optional, cast
 
 import torch
 
+from torchsparsegradutils._compat import linalg_solve_triangular_compat
 from torchsparsegradutils.utils import convert_coo_to_csr, sparse_block_diag, sparse_block_diag_split, stack_csr
 
 
@@ -154,7 +155,7 @@ class SparseTriangularSolve(torch.autograd.Function):
     See Also
     --------
     sparse_triangular_solve : User-facing function that calls this autograd function.
-    torch.triangular_solve : PyTorch's native triangular solver.
+    torch.linalg.solve_triangular : PyTorch's native dense triangular solver.
     """
 
     @staticmethod
@@ -177,11 +178,9 @@ class SparseTriangularSolve(torch.autograd.Function):
             A = convert_coo_to_csr(A)  # NOTE: triangular solve doesn't work with sparse coo
             ctx.csr = False
 
-        # NOTE: DEPRECATED: Check if a workaround for https://github.com/pytorch/pytorch/issues/88890 is needed
-
-        x = torch.triangular_solve(
-            B.detach(), A.detach(), upper=upper, unitriangular=unitriangular, transpose=transpose
-        ).solution
+        x = linalg_solve_triangular_compat(
+            A.detach(), B.detach(), upper=upper, unitriangular=unitriangular, transpose=transpose
+        )
 
         x.requires_grad = grad_flag
         ctx.save_for_backward(A, x.detach())
@@ -199,11 +198,10 @@ class SparseTriangularSolve(torch.autograd.Function):
         A, x = ctx.saved_tensors
 
         # Backprop rule: gradB = A^{-T} grad
-        # NOTE: DEPRECATED: Check if a workaround for https://github.com/pytorch/pytorch/issues/88890 is needed
 
-        gradB = torch.triangular_solve(
-            grad, A, upper=ctx.upper, transpose=not ctx.transpose, unitriangular=ctx.unitriangular
-        ).solution
+        gradB = linalg_solve_triangular_compat(
+            A, grad, upper=ctx.upper, transpose=not ctx.transpose, unitriangular=ctx.unitriangular
+        )
 
         # The gradient with respect to the matrix A seen as a dense matrix would
         # lead to a backprop rule as follows

torchsparsegradutils-0.2.6/torchsparsegradutils/tests/test_deprecated_torch_apis.py

@@ -0,0 +1,62 @@
+"""Regression checks for deprecated upstream PyTorch APIs."""
+
+import ast
+import re
+from pathlib import Path
+
+REPOSITORY_ROOT = Path(__file__).parents[2]
+PACKAGE_ROOT = REPOSITORY_ROOT / "torchsparsegradutils"
+COMPATIBILITY_MODULE = Path("_compat.py")
+DEPRECATED_DOCUMENTATION_NORM = re.compile(r"(?<![\w.])torch\.norm\s*\(")
+
+
+def attribute_name(node):
+    parts = []
+    while isinstance(node, ast.Attribute):
+        parts.append(node.attr)
+        node = node.value
+    if isinstance(node, ast.Name):
+        parts.append(node.id)
+    return ".".join(reversed(parts))
+
+
+def test_deprecated_torch_calls_are_not_reintroduced():
+    violations = []
+    compatibility_calls = []
+
+    for path in PACKAGE_ROOT.rglob("*.py"):
+        relative_path = path.relative_to(PACKAGE_ROOT)
+        tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
+
+        for node in ast.walk(tree):
+            location = f"{relative_path}:{getattr(node, 'lineno', '?')}"
+            if isinstance(node, ast.Attribute) and attribute_name(node) == "torch.jit.script":
+                violations.append(location)
+
+            if not isinstance(node, ast.Call) or not isinstance(node.func, ast.Attribute):
+                continue
+
+            call_name = attribute_name(node.func)
+            if call_name == "torch.norm":
+                violations.append(location)
+            elif call_name == "torch.triangular_solve":
+                if relative_path == COMPATIBILITY_MODULE:
+                    compatibility_calls.append(location)
+                else:
+                    violations.append(location)
+
+    assert not violations, "Deprecated PyTorch APIs found:\n" + "\n".join(violations)
+    assert len(compatibility_calls) == 1, "Expected one isolated sparse triangular compatibility call"
+
+
+def test_documentation_does_not_use_deprecated_norm_apis():
+    violations = []
+    documentation_paths = [REPOSITORY_ROOT / "README.md"]
+    documentation_paths.extend((REPOSITORY_ROOT / "docs").rglob("*.md"))
+    documentation_paths.extend((REPOSITORY_ROOT / "docs").rglob("*.rst"))
+
+    for path in documentation_paths:
+        if DEPRECATED_DOCUMENTATION_NORM.search(path.read_text(encoding="utf-8")):
+            violations.append(str(path.relative_to(REPOSITORY_ROOT)))
+
+    assert not violations, "Deprecated norm examples found:\n" + "\n".join(violations)