torchsparsegradutils 0.2.1__tar.gz → 0.2.2__tar.gz

{torchsparsegradutils-0.2.1/torchsparsegradutils.egg-info → torchsparsegradutils-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torchsparsegradutils
-Version: 0.2.1
+Version: 0.2.2
 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
@@ -19,9 +19,14 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: torch>=2.5
 Requires-Dist: scipy
-Provides-Extra: extras
-Requires-Dist: jax; extra == "extras"
-Requires-Dist: cupy; extra == "extras"
+Requires-Dist: packaging
+Provides-Extra: cupy
+Requires-Dist: cupy-cuda12x>=13.0; extra == "cupy"
+Provides-Extra: jax
+Requires-Dist: jax[cuda12]; extra == "jax"
+Provides-Extra: all
+Requires-Dist: cupy-cuda12x>=13.0; extra == "all"
+Requires-Dist: jax[cuda12]; extra == "all"
 Provides-Extra: docs
 Requires-Dist: sphinx>=7.0.0; extra == "docs"
 Requires-Dist: sphinx-rtd-theme>=1.3.0; extra == "docs"
@@ -42,7 +47,7 @@ Dynamic: license-file
 
 # torchsparsegradutils: Sparsity-preserving gradient utility tools for PyTorch
 
-[![PyPI](https://img.shields.io/pypi/v/torchsparsegradutils.svg)](https://pypi.org/project/torchsparsegradutils/) [![Python Versions](https://img.shields.io/pypi/pyversions/torchsparsegradutils.svg)](https://pypi.org/project/torchsparsegradutils/) [![Downloads](https://img.shields.io/pypi/dm/torchsparsegradutils.svg)](https://pypi.org/project/torchsparsegradutils/) ![PyTorch 2.5+](https://img.shields.io/badge/PyTorch-2.5%2B-ee4c2c?logo=pytorch) ![Tested 2.5 / 2.9 / nightly](https://img.shields.io/badge/Tested-2.5%20|%202.9%20|%20nightly-ee4c2c?logo=pytorch) [![Build](https://github.com/cai4cai/torchsparsegradutils/actions/workflows/python-package.yml/badge.svg)](https://github.com/cai4cai/torchsparsegradutils/actions/workflows/python-package.yml) [![Docs](https://readthedocs.org/projects/torchsparsegradutils/badge/?version=latest)](https://readthedocs.org/projects/torchsparsegradutils) [![Code Style: Black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black) [![License](https://img.shields.io/github/license/cai4cai/torchsparsegradutils)](LICENSE) [![status](https://joss.theoj.org/papers/6da0e92488d06f70c0a03d0a7cbfba7d/status.svg)](https://joss.theoj.org/papers/6da0e92488d06f70c0a03d0a7cbfba7d)
+[![PyPI](https://img.shields.io/pypi/v/torchsparsegradutils.svg)](https://pypi.org/project/torchsparsegradutils/) [![Python Versions](https://img.shields.io/pypi/pyversions/torchsparsegradutils.svg)](https://pypi.org/project/torchsparsegradutils/) [![Downloads](https://img.shields.io/pypi/dm/torchsparsegradutils.svg)](https://pypi.org/project/torchsparsegradutils/) ![PyTorch 2.5+](https://img.shields.io/badge/PyTorch-2.5%2B-ee4c2c?logo=pytorch) ![Tested 2.5 / 2.11 / nightly](https://img.shields.io/badge/Tested-2.5%20|%202.11%20|%20nightly-ee4c2c?logo=pytorch) [![Build](https://github.com/cai4cai/torchsparsegradutils/actions/workflows/python-package.yml/badge.svg)](https://github.com/cai4cai/torchsparsegradutils/actions/workflows/python-package.yml) [![Docs](https://readthedocs.org/projects/torchsparsegradutils/badge/?version=latest)](https://readthedocs.org/projects/torchsparsegradutils) [![Code Style: Black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black) [![License](https://img.shields.io/github/license/cai4cai/torchsparsegradutils)](LICENSE) [![status](https://joss.theoj.org/papers/6da0e92488d06f70c0a03d0a7cbfba7d/status.svg)](https://joss.theoj.org/papers/6da0e92488d06f70c0a03d0a7cbfba7d)
 
 A comprehensive collection of utility functions to work with PyTorch sparse tensors, ensuring memory efficiency and supporting various sparsity-preserving tensor operations with automatic differentiation. This package addresses fundamental gaps in PyTorch's sparse tensor ecosystem, providing essential operations that preserve sparsity in gradients during backpropagation.
 
@@ -133,17 +138,21 @@ pip install git+https://github.com/cai4cai/torchsparsegradutils
 For full functionality, install optional dependencies:
 
 ```bash
-# For CuPy sparse solver support (GPU acceleration)
-pip install cupy-cuda12x  # Replace with your CUDA version
+# For CuPy sparse solver support (GPU acceleration, requires CUDA 12.x)
+pip install torchsparsegradutils[cupy]
 
 # For JAX sparse solver support
-pip install "jax[cpu]"     # CPU version
-pip install "jax[cuda12]"  # GPU version (replace with your CUDA version)
+pip install torchsparsegradutils[jax]
+
+# Install all optional dependencies
+pip install torchsparsegradutils[all]
 
 # For benchmarking and testing
 pip install scipy matplotlib pandas tqdm pytest
 ```
 
+> **Note:** The CuPy extra installs `cupy-cuda12x>=13.0`. If you are using a different CUDA version, install the appropriate CuPy package manually (e.g. `pip install cupy-cuda11x`).
+
 ### Requirements
 
 - **Python**: ≥ 3.10
@@ -160,7 +169,7 @@ Our comprehensive benchmark suite demonstrates significant performance improveme
 
 ![Sparse Triangular Solve Suite Performance (int32/float32 COO)](torchsparsegradutils/benchmarks/benchmark_visualizations/triangular_solve_suitesparse_performance_int32_float32_coo.png)
 
-![Sparse Genertic Solve Suite Performance (int32/float32 COO)](torchsparsegradutils/benchmarks/benchmark_visualizations/sparse_solve_suite_performance_int32_float32_coo.png)
+![Sparse Generic Solve Suite Performance (int32/float32 COO)](torchsparsegradutils/benchmarks/benchmark_visualizations/sparse_solve_suite_performance_int32_float32_coo.png)
 
 ## 🚀 Quick Start
 

{torchsparsegradutils-0.2.1 → torchsparsegradutils-0.2.2}/README.md

@@ -1,6 +1,6 @@
 # torchsparsegradutils: Sparsity-preserving gradient utility tools for PyTorch
 
-[![PyPI](https://img.shields.io/pypi/v/torchsparsegradutils.svg)](https://pypi.org/project/torchsparsegradutils/) [![Python Versions](https://img.shields.io/pypi/pyversions/torchsparsegradutils.svg)](https://pypi.org/project/torchsparsegradutils/) [![Downloads](https://img.shields.io/pypi/dm/torchsparsegradutils.svg)](https://pypi.org/project/torchsparsegradutils/) ![PyTorch 2.5+](https://img.shields.io/badge/PyTorch-2.5%2B-ee4c2c?logo=pytorch) ![Tested 2.5 / 2.9 / nightly](https://img.shields.io/badge/Tested-2.5%20|%202.9%20|%20nightly-ee4c2c?logo=pytorch) [![Build](https://github.com/cai4cai/torchsparsegradutils/actions/workflows/python-package.yml/badge.svg)](https://github.com/cai4cai/torchsparsegradutils/actions/workflows/python-package.yml) [![Docs](https://readthedocs.org/projects/torchsparsegradutils/badge/?version=latest)](https://readthedocs.org/projects/torchsparsegradutils) [![Code Style: Black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black) [![License](https://img.shields.io/github/license/cai4cai/torchsparsegradutils)](LICENSE) [![status](https://joss.theoj.org/papers/6da0e92488d06f70c0a03d0a7cbfba7d/status.svg)](https://joss.theoj.org/papers/6da0e92488d06f70c0a03d0a7cbfba7d)
+[![PyPI](https://img.shields.io/pypi/v/torchsparsegradutils.svg)](https://pypi.org/project/torchsparsegradutils/) [![Python Versions](https://img.shields.io/pypi/pyversions/torchsparsegradutils.svg)](https://pypi.org/project/torchsparsegradutils/) [![Downloads](https://img.shields.io/pypi/dm/torchsparsegradutils.svg)](https://pypi.org/project/torchsparsegradutils/) ![PyTorch 2.5+](https://img.shields.io/badge/PyTorch-2.5%2B-ee4c2c?logo=pytorch) ![Tested 2.5 / 2.11 / nightly](https://img.shields.io/badge/Tested-2.5%20|%202.11%20|%20nightly-ee4c2c?logo=pytorch) [![Build](https://github.com/cai4cai/torchsparsegradutils/actions/workflows/python-package.yml/badge.svg)](https://github.com/cai4cai/torchsparsegradutils/actions/workflows/python-package.yml) [![Docs](https://readthedocs.org/projects/torchsparsegradutils/badge/?version=latest)](https://readthedocs.org/projects/torchsparsegradutils) [![Code Style: Black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black) [![License](https://img.shields.io/github/license/cai4cai/torchsparsegradutils)](LICENSE) [![status](https://joss.theoj.org/papers/6da0e92488d06f70c0a03d0a7cbfba7d/status.svg)](https://joss.theoj.org/papers/6da0e92488d06f70c0a03d0a7cbfba7d)
 
 A comprehensive collection of utility functions to work with PyTorch sparse tensors, ensuring memory efficiency and supporting various sparsity-preserving tensor operations with automatic differentiation. This package addresses fundamental gaps in PyTorch's sparse tensor ecosystem, providing essential operations that preserve sparsity in gradients during backpropagation.
 
@@ -91,17 +91,21 @@ pip install git+https://github.com/cai4cai/torchsparsegradutils
 For full functionality, install optional dependencies:
 
 ```bash
-# For CuPy sparse solver support (GPU acceleration)
-pip install cupy-cuda12x  # Replace with your CUDA version
+# For CuPy sparse solver support (GPU acceleration, requires CUDA 12.x)
+pip install torchsparsegradutils[cupy]
 
 # For JAX sparse solver support
-pip install "jax[cpu]"     # CPU version
-pip install "jax[cuda12]"  # GPU version (replace with your CUDA version)
+pip install torchsparsegradutils[jax]
+
+# Install all optional dependencies
+pip install torchsparsegradutils[all]
 
 # For benchmarking and testing
 pip install scipy matplotlib pandas tqdm pytest
 ```
 
+> **Note:** The CuPy extra installs `cupy-cuda12x>=13.0`. If you are using a different CUDA version, install the appropriate CuPy package manually (e.g. `pip install cupy-cuda11x`).
+
 ### Requirements
 
 - **Python**: ≥ 3.10
@@ -118,7 +122,7 @@ Our comprehensive benchmark suite demonstrates significant performance improveme
 
 ![Sparse Triangular Solve Suite Performance (int32/float32 COO)](torchsparsegradutils/benchmarks/benchmark_visualizations/triangular_solve_suitesparse_performance_int32_float32_coo.png)
 
-![Sparse Genertic Solve Suite Performance (int32/float32 COO)](torchsparsegradutils/benchmarks/benchmark_visualizations/sparse_solve_suite_performance_int32_float32_coo.png)
+![Sparse Generic Solve Suite Performance (int32/float32 COO)](torchsparsegradutils/benchmarks/benchmark_visualizations/sparse_solve_suite_performance_int32_float32_coo.png)
 
 ## 🚀 Quick Start
 
@@ -658,4 +662,4 @@ dist_stable = SparseMultivariateNormal(
 - **No SPD Constraints**: Doesn't require strict positive definiteness
 - **Better Conditioning**: Diagonal component can be controlled independently
 
-**Status**: This is a known limitation of the LL^T precision formulation. LDL^T parameterization is the recommended approach for precision matrices.
+**Status**: This is a known limitation of the LL^T precision formulation. LDL^T parameterization is the recommended approach for precision matrices.

{torchsparsegradutils-0.2.1 → torchsparsegradutils-0.2.2}/docs/source/conf.py

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

{torchsparsegradutils-0.2.1 → torchsparsegradutils-0.2.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "torchsparsegradutils"
-version = "0.2.1"
+version = "0.2.2"
 description = "A collection of utility functions to work with PyTorch sparse tensors"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -22,10 +22,13 @@ classifiers = [
 dependencies = [
     "torch>=2.5",
     "scipy",
+    "packaging",
 ]
 
 [project.optional-dependencies]
-extras = ["jax", "cupy"]
+cupy = ["cupy-cuda12x>=13.0"]
+jax = ["jax[cuda12]"]
+all = ["cupy-cuda12x>=13.0", "jax[cuda12]"]
 docs = [
     "sphinx>=7.0.0",
     "sphinx-rtd-theme>=1.3.0",
@@ -99,6 +102,9 @@ addopts = ["-v"]
 testpaths = [
     "torchsparsegradutils/tests",
 ]
+pythonpath = [
+    "torchsparsegradutils/tests",
+]
 doctest_optionflags = [
     "NORMALIZE_WHITESPACE",
     "IGNORE_EXCEPTION_DETAIL",

{torchsparsegradutils-0.2.1 → torchsparsegradutils-0.2.2}/torchsparsegradutils/cupy/__init__.py

@@ -11,7 +11,16 @@ if cupy_spec is None:
 else:
     have_cupy = True
 
-from .cupy_bindings import _get_array_modules, c2t_coo, c2t_csr, t2c_coo, t2c_csr
+from .cupy_bindings import _backend_to_torch, _get_array_modules, _torch_to_backend, c2t_coo, c2t_csr, t2c_coo, t2c_csr
 from .cupy_sparse_solve import sparse_solve_c4t
 
-__all__ = ["c2t_coo", "t2c_coo", "c2t_csr", "t2c_csr", "_get_array_modules", "sparse_solve_c4t"]
+__all__ = [
+    "c2t_coo",
+    "t2c_coo",
+    "c2t_csr",
+    "t2c_csr",
+    "_get_array_modules",
+    "_torch_to_backend",
+    "_backend_to_torch",
+    "sparse_solve_c4t",
+]

{torchsparsegradutils-0.2.1 → torchsparsegradutils-0.2.2}/torchsparsegradutils/cupy/cupy_bindings.py

@@ -38,6 +38,20 @@ import scipy.sparse.linalg
 import torch
 
 
+def _torch_to_backend(tensor: torch.Tensor, xp) -> Any:
+    r"""Convert a PyTorch tensor to a backend array (NumPy or CuPy) using DLPack for CUDA."""
+    if tsgucupy.have_cupy and xp is cp:
+        return cp.from_dlpack(tensor)
+    return xp.asarray(tensor)
+
+
+def _backend_to_torch(array) -> torch.Tensor:
+    r"""Convert a backend array (NumPy or CuPy) to a PyTorch tensor using DLPack for CUDA."""
+    if tsgucupy.have_cupy and isinstance(array, cp.ndarray):
+        return torch.from_dlpack(array)
+    return torch.as_tensor(array)
+
+
 def _get_array_modules(x: Any) -> Tuple[Any, Any]:
     r"""
     Select dense & sparse array modules (NumPy/SciPy or CuPy/cupyx.scipy.sparse).
@@ -116,9 +130,9 @@ def t2c_csr(x_torch: torch.Tensor) -> Any:
     """
     xp, xsp = _get_array_modules(x_torch)
 
-    data_c = xp.asarray(x_torch.values())
-    col_idx_c = xp.asarray(x_torch.col_indices())
-    ind_ptr_c = xp.asarray(x_torch.crow_indices())
+    data_c = _torch_to_backend(x_torch.values(), xp)
+    col_idx_c = _torch_to_backend(x_torch.col_indices(), xp)
+    ind_ptr_c = _torch_to_backend(x_torch.crow_indices(), xp)
     x_cupy = xsp.csr_matrix((data_c, col_idx_c, ind_ptr_c), shape=x_torch.shape)
     return x_cupy
 
@@ -150,9 +164,9 @@ def c2t_csr(x_cupy: Any) -> torch.Tensor:
     >>> x.layout is torch.sparse_csr
     True
     """
-    data_t = torch.as_tensor(x_cupy.data)
-    idices_t = torch.as_tensor(x_cupy.indices)
-    ind_ptr_t = torch.as_tensor(x_cupy.indptr)
+    data_t = _backend_to_torch(x_cupy.data)
+    idices_t = _backend_to_torch(x_cupy.indices)
+    ind_ptr_t = _backend_to_torch(x_cupy.indptr)
     x_torch = torch.sparse_csr_tensor(ind_ptr_t, idices_t, data_t, x_cupy.shape)
     return x_torch
 
@@ -205,8 +219,8 @@ def t2c_coo(x_torch: torch.Tensor) -> Any:
             "Requested a conversion from torch to cupy/numpy on a non-coalesced tensor -> coalescing implicitly"
         )
         x_torch = x_torch.coalesce()
-    data_c = xp.asarray(x_torch.values())
-    idx_cp = xp.asarray(x_torch.indices())
+    data_c = _torch_to_backend(x_torch.values(), xp)
+    idx_cp = _torch_to_backend(x_torch.indices(), xp)
     x_cupy = xsp.coo_matrix((data_c, idx_cp), shape=x_torch.shape)
     return x_cupy
 
@@ -238,8 +252,8 @@ def c2t_coo(x_cupy: Any) -> torch.Tensor:
     >>> x.layout is torch.sparse_coo
     True
     """
-    data_t = torch.as_tensor(x_cupy.data)
-    row_t = torch.as_tensor(x_cupy.row)
-    col_t = torch.as_tensor(x_cupy.col)
+    data_t = _backend_to_torch(x_cupy.data)
+    row_t = _backend_to_torch(x_cupy.row)
+    col_t = _backend_to_torch(x_cupy.col)
     x_torch = torch.sparse_coo_tensor(torch.stack([row_t, col_t], dim=0), data_t, x_cupy.shape)
     return x_torch

{torchsparsegradutils-0.2.1 → torchsparsegradutils-0.2.2}/torchsparsegradutils/cupy/cupy_sparse_solve.py

@@ -1,3 +1,4 @@
+import inspect
 import warnings
 from typing import Any, Callable, Dict, Optional, Tuple, Union
 
@@ -8,6 +9,90 @@ import torchsparsegradutils.cupy as tsgucupy
 # from cupyx.scipy.sparse.linalg import cg, cgs, minres, gmres, spsolve
 
 
+def _wrap_iterative_solver(base_solver, backend_type, solver_name=None):
+    """Wrap an iterative solver to handle parameter mapping and return format."""
+
+    # Introspect the actual solver signature to determine accepted parameter names
+    try:
+        sig_params = set(inspect.signature(base_solver).parameters.keys())
+    except (ValueError, TypeError):
+        sig_params = set()
+
+    def wrapped_solver(A, b, **solver_kwargs):
+        # Create a copy to avoid modifying the original
+        filtered_kwargs = solver_kwargs.copy()
+
+        # Extract tolerance parameter and map to correct name for the backend
+        tolerance = filtered_kwargs.pop("tol", None)
+        atol = filtered_kwargs.pop("atol", None)
+
+        if tolerance is not None:
+            # Use introspection to determine correct tolerance parameter name
+            if "rtol" in sig_params:
+                filtered_kwargs["rtol"] = tolerance
+            elif "tol" in sig_params:
+                filtered_kwargs["tol"] = tolerance
+
+        # Handle atol parameter if the solver accepts it
+        if atol is not None and "atol" in sig_params:
+            filtered_kwargs["atol"] = atol
+
+        # Filter to only parameters the solver actually accepts
+        common_params = {"x0", "M", "callback", "show", "check"}
+        if sig_params:
+            final_kwargs = {k: v for k, v in filtered_kwargs.items() if k in sig_params or k in common_params}
+        else:
+            # Fallback: pass everything if we couldn't introspect
+            final_kwargs = filtered_kwargs
+
+        # Call the base solver
+        result = base_solver(A, b, **final_kwargs)
+
+        # Handle return format - some solvers return (solution, info) tuples
+        if isinstance(result, tuple):
+            return result[0]  # Return just the solution
+        return result
+
+    return wrapped_solver
+
+
+def _wrap_direct_solver(base_solver):
+    """Wrap a direct solver to ignore tolerance parameters."""
+
+    def wrapped_solver(A, b, **solver_kwargs):
+        # Direct solvers don't use iterative solver parameters, so ignore them all
+        filtered_kwargs = {
+            k: v
+            for k, v in solver_kwargs.items()
+            if k not in ["tol", "tolerance", "atol", "rtol", "maxiter", "matvec_max"]
+        }
+        return base_solver(A, b, **filtered_kwargs)
+
+    return wrapped_solver
+
+
+def _get_solver_function(solver_name, xsp, device):
+    """Get the appropriate solver function based on the backend."""
+    if solver_name is None or callable(solver_name):
+        return solver_name
+
+    # Determine backend type
+    backend_type = "scipy" if device.type == "cpu" else "cupy"
+
+    solver_map = {
+        "cg": _wrap_iterative_solver(xsp.linalg.cg, backend_type, "cg"),
+        "cgs": _wrap_iterative_solver(xsp.linalg.cgs, backend_type, "cgs"),
+        "minres": _wrap_iterative_solver(xsp.linalg.minres, backend_type, "minres"),
+        "gmres": _wrap_iterative_solver(xsp.linalg.gmres, backend_type, "gmres"),
+        "spsolve": _wrap_direct_solver(xsp.linalg.spsolve),
+    }
+
+    if solver_name not in solver_map:
+        raise ValueError(f"Unknown solver: {solver_name}. Supported solvers: {list(solver_map.keys())}")
+
+    return solver_map[solver_name]
+
+
 def sparse_solve_c4t(
     A: torch.Tensor,
     B: torch.Tensor,
@@ -154,94 +239,7 @@ def sparse_solve_c4t(
             f"Use transpose_solve='spsolve' or transpose_solve=None for multi-RHS problems."
         )
 
-    # Convert string solver names to actual solver functions
-    def _get_solver_function(solver_name, xsp):
-        """Get the appropriate solver function based on the backend."""
-        if solver_name is None or callable(solver_name):
-            return solver_name
-
-        def _wrap_iterative_solver(base_solver, backend_type, solver_name=None):
-            """Wrap an iterative solver to handle parameter mapping and return format."""
-
-            def wrapped_solver(A, b, **solver_kwargs):
-                # Create a copy to avoid modifying the original
-                filtered_kwargs = solver_kwargs.copy()
-
-                # Extract tolerance parameter and map to correct name for the backend
-                tolerance = filtered_kwargs.pop("tol", None)
-                atol = filtered_kwargs.pop("atol", None)
-
-                # Define solver-specific parameter support
-                solver_params = {
-                    "cg": {"rtol", "maxiter", "atol"} if backend_type == "scipy" else {"tol", "maxiter", "atol"},
-                    "cgs": {"rtol", "maxiter", "atol"} if backend_type == "scipy" else {"tol", "maxiter", "atol"},
-                    "minres": {"rtol", "maxiter", "shift"} if backend_type == "scipy" else {"tol", "maxiter"},
-                    "gmres": {"rtol", "maxiter", "atol"} if backend_type == "scipy" else {"tol", "maxiter", "atol"},
-                }
-
-                # Get supported parameters for this solver
-                supported_params = solver_params.get(solver_name, set())
-
-                if tolerance is not None:
-                    # Map tolerance parameter based on backend and solver support
-                    if backend_type == "scipy" and "rtol" in supported_params:
-                        filtered_kwargs["rtol"] = tolerance
-                    elif backend_type == "cupy" and "tol" in supported_params:
-                        filtered_kwargs["tol"] = tolerance
-
-                # Handle atol parameter if supported
-                if atol is not None and "atol" in supported_params:
-                    filtered_kwargs["atol"] = atol
-
-                # Filter out unsupported parameters
-                final_kwargs = {
-                    k: v
-                    for k, v in filtered_kwargs.items()
-                    if k in supported_params
-                    or k in {"x0", "M", "callback", "show", "check"}  # Always allow common parameters
-                }
-
-                # Call the base solver
-                result = base_solver(A, b, **final_kwargs)
-
-                # Handle return format - some solvers return (solution, info) tuples
-                if isinstance(result, tuple):
-                    return result[0]  # Return just the solution
-                return result
-
-            return wrapped_solver
-
-        def _wrap_direct_solver(base_solver):
-            """Wrap a direct solver to ignore tolerance parameters."""
-
-            def wrapped_solver(A, b, **solver_kwargs):
-                # Direct solvers don't use iterative solver parameters, so ignore them all
-                filtered_kwargs = {
-                    k: v
-                    for k, v in solver_kwargs.items()
-                    if k not in ["tol", "tolerance", "atol", "rtol", "maxiter", "matvec_max"]
-                }
-                return base_solver(A, b, **filtered_kwargs)
-
-            return wrapped_solver
-
-        # Determine backend type
-        backend_type = "scipy" if A.device.type == "cpu" else "cupy"
-
-        solver_map = {
-            "cg": _wrap_iterative_solver(xsp.linalg.cg, backend_type, "cg"),
-            "cgs": _wrap_iterative_solver(xsp.linalg.cgs, backend_type, "cgs"),
-            "minres": _wrap_iterative_solver(xsp.linalg.minres, backend_type, "minres"),
-            "gmres": _wrap_iterative_solver(xsp.linalg.gmres, backend_type, "gmres"),
-            "spsolve": _wrap_direct_solver(xsp.linalg.spsolve),
-        }
-
-        if solver_name not in solver_map:
-            raise ValueError(f"Unknown solver: {solver_name}. Supported solvers: {list(solver_map.keys())}")
-
-        return solver_map[solver_name]
-
-    # Get the appropriate backend modules
+    # Convert string solver names to functions
     xp, xsp = tsgucupy._get_array_modules(A.data)
 
     # Warn about dtype issues with minres on CPU
@@ -262,9 +260,8 @@ def sparse_solve_c4t(
             stacklevel=2,
         )
 
-    # Convert string solver names to functions
-    solve_func = _get_solver_function(solve, xsp)
-    transpose_solve_func = _get_solver_function(transpose_solve, xsp)
+    solve_func = _get_solver_function(solve, xsp, A.device)
+    transpose_solve_func = _get_solver_function(transpose_solve, xsp, A.device)
 
     return SparseSolveC4T.apply(A, B, solve_func, transpose_solve_func, kwargs)
 
@@ -307,7 +304,7 @@ class SparseSolveC4T(torch.autograd.Function):
             A_c = tsgucupy.t2c_csr(A.detach())
         else:
             raise TypeError(f"Unsupported layout type: {A.layout}")
-        B_c = xp.asarray(B.detach())
+        B_c = tsgucupy._torch_to_backend(B.detach(), xp)
 
         # Solve the sparse system
         ctx.factorisedsolver = None
@@ -329,7 +326,7 @@ class SparseSolveC4T(torch.autograd.Function):
             # If the solver returns a tuple, we assume the first element is the solution
             x_c = x_c[0]
 
-        x = torch.as_tensor(x_c, device=A.device)
+        x = tsgucupy._backend_to_torch(x_c)
 
         # Ensure output dtype matches input dtype
         if x.dtype != A.dtype:
@@ -354,7 +351,7 @@ class SparseSolveC4T(torch.autograd.Function):
             x = x.unsqueeze(-1)
             grad = grad.unsqueeze(-1)
 
-        grad_c = xp.asarray(grad.detach())
+        grad_c = tsgucupy._torch_to_backend(grad.detach(), xp)
 
         # Backprop rule: gradB = A^{-T} grad
         if ctx.transpose_solve is not None:
@@ -369,7 +366,7 @@ class SparseSolveC4T(torch.autograd.Function):
             # If the solver returns a tuple, we assume the first element is the gradient
             gradB_c = gradB_c[0]
 
-        gradB = torch.as_tensor(gradB_c, device=A.device)
+        gradB = tsgucupy._backend_to_torch(gradB_c)
 
         # Ensure gradient dtype matches input dtype
         if gradB.dtype != A.dtype:

{torchsparsegradutils-0.2.1 → torchsparsegradutils-0.2.2}/torchsparsegradutils/indexed_matmul.py

@@ -1,4 +1,5 @@
 import torch
+from packaging.version import parse as parse_version
 
 try:
     import dgl.ops as dglops
@@ -73,7 +74,7 @@ def segment_mm(a: torch.Tensor, b: torch.Tensor, seglen_a: torch.Tensor) -> torc
         >>> segment_mm(a, b, seglen_a).shape
         torch.Size([18, 2])
     """
-    if torch.__version__ < (2, 4):
+    if parse_version(torch.__version__) < parse_version("2.4"):
         raise NotImplementedError("PyTorch version is too old for nested tensors")
 
     if dgl_installed:
@@ -177,7 +178,7 @@ def gather_mm(a: torch.Tensor, b: torch.Tensor, idx_b: torch.Tensor) -> torch.Te
         tensor([[1., 2.],
                 [6., 8.]])
     """
-    if torch.__version__ < (2, 4):
+    if parse_version(torch.__version__) < parse_version("2.4"):
         raise NotImplementedError("PyTorch version is too old for nested tensors")
 
     if dgl_installed:

{torchsparsegradutils-0.2.1 → torchsparsegradutils-0.2.2}/torchsparsegradutils/sparse_matmul.py

@@ -150,7 +150,7 @@ class SparseMatMul(torch.autograd.Function):
 
         if ctx.batch_size is not None:
             A = sparse_block_diag(*A)
-            B = torch.cat([*B])
+            B = B.reshape(-1, B.size(-1))
 
         x = torch.sparse.mm(A, B)
 
@@ -194,7 +194,7 @@ class SparseMatMul(torch.autograd.Function):
                 raise ValueError(f"Unsupported layout: {A.layout}")
 
             if ctx.batch_size is not None:
-                grad_for_A = torch.cat([*grad])
+                grad_for_A = grad.reshape(-1, grad.size(-1))
             else:
                 grad_for_A = grad
 
@@ -221,7 +221,7 @@ class SparseMatMul(torch.autograd.Function):
         # -------- Only compute gradB if needed --------
         if ctx.needs_input_grad[1]:
             if ctx.batch_size is not None:
-                grad_for_B = torch.cat([*grad])
+                grad_for_B = grad.reshape(-1, grad.size(-1))
             else:
                 grad_for_B = grad
 

{torchsparsegradutils-0.2.1 → torchsparsegradutils-0.2.2}/torchsparsegradutils/sparse_solve.py

@@ -171,7 +171,7 @@ class SparseTriangularSolve(torch.autograd.Function):
 
         if ctx.batch_size is not None:
             A = sparse_block_diag(*A)
-            B = torch.cat([*B])
+            B = B.reshape(-1, B.size(-1))
 
         if A.layout == torch.sparse_coo:
             A = convert_coo_to_csr(A)  # NOTE: triangular solve doesn't work with sparse coo
@@ -194,7 +194,7 @@ class SparseTriangularSolve(torch.autograd.Function):
     @staticmethod
     def backward(ctx, grad):  # type: ignore[override]
         if ctx.batch_size is not None:
-            grad = torch.cat([*grad])
+            grad = grad.reshape(-1, grad.size(-1))
 
         A, x = ctx.saved_tensors
 

torchsparsegradutils-0.2.2/torchsparsegradutils/tests/conftest.py

@@ -0,0 +1,48 @@
+import os
+import random
+
+import numpy as np
+import pytest
+import torch
+
+SEED = 42
+
+
+def _seed_unlocked() -> bool:
+    unlock = os.getenv("TSGU_UNLOCK_SEED", os.getenv("UNLOCK_SEED", "false"))
+    return unlock.lower() in {"1", "true", "yes", "on"}
+
+
+def _seed_all(seed: int = SEED) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def pytest_configure(config):
+    if not _seed_unlocked():
+        _seed_all()
+
+
+@pytest.fixture(autouse=True)
+def seed_rng():
+    if _seed_unlocked():
+        yield
+        return
+
+    torch_state = torch.get_rng_state()
+    cuda_states = torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None
+    random_state = random.getstate()
+    numpy_state = np.random.get_state()
+
+    _seed_all()
+    try:
+        yield
+    finally:
+        torch.set_rng_state(torch_state)
+        if cuda_states is not None:
+            torch.cuda.set_rng_state_all(cuda_states)
+        random.setstate(random_state)
+        np.random.set_state(numpy_state)