sparse-kappa 0.0.2__tar.gz → 0.0.3__tar.gz

{sparse_kappa-0.0.2/sparse_kappa.egg-info → sparse_kappa-0.0.3}/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: sparse-kappa
-Version: 0.0.2
-Summary: GPU-accelerated sparse matrix condition number estimation using CuPy
+Version: 0.0.3
+Summary: PyTorch-accelerated sparse matrix condition number estimation
 Home-page: https://github.com/chenxinye/sparse-kappa
 Author: Xinye Chen
 Author-email: Xinye Chen <xinyechenai@gmail.com>
@@ -19,8 +19,9 @@ Classifier: Programming Language :: Python :: 3.11
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: cupy>=10.0.0
 Requires-Dist: numpy>=1.20.0
+Requires-Dist: torch>=2.0.0
+Requires-Dist: myst-parser>=2.0
 Provides-Extra: dev
 Requires-Dist: pytest>=6.0; extra == "dev"
 Requires-Dist: pytest-cov>=2.0; extra == "dev"
@@ -35,28 +36,31 @@ Dynamic: requires-python
 
 # sparse-kappa
 
-**Condition Number Estimation on GPUs for Sparse Matrices**
+[![!pypi](https://img.shields.io/pypi/v/sparse-kappa?color=indigo)](https://pypi.org/project/sparse-kappa/)
+[![License: MIT](https://anaconda.org/conda-forge/sparse-kappaclustering/badges/license.svg)]([https://github.com/nla-group/sparse-kappa/blob/master/LICENSE](https://opensource.org/licenses/MIT))
+[![Documentation Status](https://readthedocs.org/projects/sparse-kappa/badge/?version=latest)](https://sparse-kappa.readthedocs.io/en/latest/)
 
 
-
-</div>
+**Condition Number Estimation on CPUs/GPUs for Sparse Matrices**
 
 
-``sparse-kappa`` is a GPU-accelerated library for estimating condition numbers of sparse matrices using CuPy. It supports a variaty of estimation method associated with linear solvers. ``sparse-kappa`` is designed for benchmarking condition number estimator and practical use in science and engineering community.  
 
+</div>
 
 
+``sparse-kappa`` is a GPU-accelerated library for estimating condition numbers of sparse matrices using PyTorch. It supports a variaty of estimation method associated with linear solvers. ``sparse-kappa`` is designed for benchmarking condition number estimator and practical use in science and engineering community.  
 
 
 ## Features
 
-- **GPU-Accelerated**: All computations run on NVIDIA GPUs via CuPy
+- **GPU-Accelerated**: All computations run on NVIDIA GPUs via PyTorch
 - **Multiple Norms**: Support for 1-norm and 2-norm condition numbers
 - **Rich Algorithm Suite**:
   - **1-norm**: Hager-Higham, Power iteration, Oettli-Prager sampling, Block Hager
-  - **2-norm**: Power method, Lanczos, Arnoldi, Golub-Kahan bidiagonalization
-  - **CuPy integrations**: SVDS, EIGSH, LOBPCG wrappers
+  - **2-norm**: Power method, Lanczos, Golub-Kahan bidiagonalization
+  - **PyTorch integrations**: SVDS, EIGSH, LOBPCG wrappers
 - **Flexible Solver System**: LU, LSMR, CG, GMRES, Direct, Auto-selection
+- **GNN Prediction Module**: Train reusable models that predict condition numbers directly or via inverse-norm prediction
 - **Smart LU Caching**: Reuses factorizations for multiple solves (10-20x speedup)
 - **Memory Efficient**: Designed for large sparse matrices
 
@@ -71,14 +75,14 @@ pip install sparse-kappa
 
 ```bash
 git clone https://github.com/chenxinye/sparse-kappa
-pip install cupy-cuda11x  # or cupy-cuda12x for CUDA 12
+pip install torch
 pip install -e .
 ```
 
 ## Quick Start
 
 ```python
-import cupyx.scipy.sparse as sp
+from sparse_kappa.backend import sparse as sp
 from sparse_kappa import cond_estimate
 
 # Create sparse matrix
@@ -107,16 +111,16 @@ cond = cond_estimate(A, norm=1, method='hager-higham', solver='lu')
 
 ## 2-Norm Methods
 
-| Method       | Description                        | Best For                     | Complexity      |
-|--------------|------------------------------------|------------------------------|-----------------|
-| `svds`       | Partial SVD (most accurate)        | Small-medium matrices (<5k)  | O(k·nnz)        |
-| `eigsh`      | Symmetric eigenvalue solver        | Symmetric matrices           | O(k·nnz)        |
-| `lobpcg`     | Block preconditioned CG            | Large matrices               | O(k·nnz)        |
-| `power`      | Power iteration                    | Quick estimates              | O(k·nnz)        |
-| `lanczos`    | Lanczos tridiagonalization         | Medium matrices              | O(k²·nnz)       |
-| `arnoldi`    | Arnoldi iteration                  | Non-symmetric                | O(k²·nnz)       |
-| `golub-kahan`| Bidiagonalization                  | Numerically stable           | O(k·nnz)        |
-| `auto`       | Automatic selection                | All cases                    | -               |
+| Method            | Description                                                | Best For                                | Complexity      |
+|-------------------|------------------------------------------------------------|-----------------------------------------|-----------------|
+| `svds`            | Partial SVD (most accurate)                                | Small-medium matrices (<5k)             | O(k·nnz)        |
+| `eigsh`           | Symmetric eigenvalue solver                                | Symmetric / Hermitian matrices          | O(k·nnz)        |
+| `lobpcg`          | Block preconditioned CG                                    | Large matrices                          | O(k·nnz)        |
+| `power`           | Power iteration                                            | Quick estimates                         | O(k·nnz)        |
+| `lanczos`         | Lanczos tridiagonalization                                 | Medium symmetric matrices               | O(k²·nnz)       |
+| `lanczos_unsym`   | Lanczos-style condition estimation via `eigsh` on `A^H A`  | Non-symmetric / rectangular matrices    | O(k·nnz)        |
+| `golub-kahan`     | Bidiagonalization                                          | Numerically stable                      | O(k·nnz)        |
+| `auto`            | Automatic selection                                        | All cases                               | -               |
 
 ## Solver Options
 
@@ -135,12 +139,70 @@ All 1-norm methods support flexible solver selection:
 **Legend:**  
 `k` = iterations, `b` = block size, `m` = samples, `nnz` = non-zeros
 
+## GNN-Based Prediction
+
+The `sparse_kappa.gnn` module learns a mapping from sparse matrices to
+condition-number related scalars. It supports two training targets:
+
+- `target="condition"`: predict `kappa(A)` directly.
+- `target="inverse_norm"`: predict `||A^{-1}||`, then compute `||A|| * ||A^{-1}||` at inference time.
+
+The default graph builder turns a sparse matrix into a row/column bipartite
+graph. You can replace the feature extractor, model, optimizer, scheduler,
+loss function, and validation callback.
+
+```python
+from sparse_kappa import TrainingConfig, train_gnn_condition_estimator
+from sparse_kappa.gnn import GNNConditionEstimator
+from sparse_kappa.backend import sparse as sp
+
+train_samples = [
+    {"matrix": A0, "condition_number": 12.3, "norm_Ainv": 0.42},
+    {"matrix": A1, "condition_number": 18.9, "norm_Ainv": 0.61},
+]
+
+# Scheme 1: direct condition-number prediction
+config = TrainingConfig(target="condition", norm=2, epochs=100, lr=1e-3)
+estimator = train_gnn_condition_estimator(
+    train_samples,
+    val_data=None,
+    config=config,
+    save_path="models/gnn_cond.pt",
+)
+
+# Load and predict one matrix or a list of matrices.
+estimator = GNNConditionEstimator.load("models/gnn_cond.pt")
+pred = estimator.predict(sp.random(100, 100, density=0.02, format="csr"))
+
+# Scheme 2: inverse-norm prediction. At inference, sparse-kappa computes ||A||
+# for the requested norm and returns ||A|| * predicted ||A^{-1}||.
+inv_config = TrainingConfig(target="inverse_norm", norm=1, epochs=100)
+inv_estimator = train_gnn_condition_estimator(train_samples, config=inv_config)
+result = inv_estimator.predict(A_test, return_dict=True)
+print(result["condition_number"], result["norm_A"], result["norm_Ainv"])
+```
+
+Customization hooks follow the same shape:
+
+```python
+import torch
+
+estimator.fit(
+    train_samples,
+    val_data=val_samples,
+    optimizer_factory=lambda params: torch.optim.Adam(params, lr=5e-4),
+    scheduler_factory=lambda opt: torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=50),
+    loss_fn=torch.nn.SmoothL1Loss(),
+    validator=my_validation_callback,
+)
+```
+
 ## Examples
 
 ### Example 1: Compare Methods
 
 ```python
-import cupyx.scipy.sparse as sp
+from sparse_kappa.backend import sparse as sp
 from sparse_kappa import cond_estimate
 
 A = sp.random(2000, 2000, density=0.005, format='csr')
@@ -248,10 +310,13 @@ Contributions welcome! Please submit issues and pull requests on GitHub.
 If you use this library in your research, please cite:
 
 ```bibtex
-@software{sparse_kappa,
-  title={Sparse Matrices Condition Number Estimation on GPUs},
-  author={Xinye Chen},
-  year={2026},
-  url={https://github.com/chenxinye/sparse_kappa}
+@misc{carson2026estimatingconditionnumbergraph,
+      title={Estimating condition number with Graph Neural Networks}, 
+      author={Erin Carson and Xinye Chen},
+      year={2026},
+      eprint={2603.10277},
+      archivePrefix={arXiv},
+      primaryClass={cs.LG},
+      url={https://arxiv.org/abs/2603.10277}, 
 }
 ```

{sparse_kappa-0.0.2 → sparse_kappa-0.0.3}/README.md

@@ -2,28 +2,31 @@
 
 # sparse-kappa
 
-**Condition Number Estimation on GPUs for Sparse Matrices**
+[![!pypi](https://img.shields.io/pypi/v/sparse-kappa?color=indigo)](https://pypi.org/project/sparse-kappa/)
+[![License: MIT](https://anaconda.org/conda-forge/sparse-kappaclustering/badges/license.svg)]([https://github.com/nla-group/sparse-kappa/blob/master/LICENSE](https://opensource.org/licenses/MIT))
+[![Documentation Status](https://readthedocs.org/projects/sparse-kappa/badge/?version=latest)](https://sparse-kappa.readthedocs.io/en/latest/)
 
 
-
-</div>
+**Condition Number Estimation on CPUs/GPUs for Sparse Matrices**
 
 
-``sparse-kappa`` is a GPU-accelerated library for estimating condition numbers of sparse matrices using CuPy. It supports a variaty of estimation method associated with linear solvers. ``sparse-kappa`` is designed for benchmarking condition number estimator and practical use in science and engineering community.  
 
+</div>
 
 
+``sparse-kappa`` is a GPU-accelerated library for estimating condition numbers of sparse matrices using PyTorch. It supports a variaty of estimation method associated with linear solvers. ``sparse-kappa`` is designed for benchmarking condition number estimator and practical use in science and engineering community.  
 
 
 ## Features
 
-- **GPU-Accelerated**: All computations run on NVIDIA GPUs via CuPy
+- **GPU-Accelerated**: All computations run on NVIDIA GPUs via PyTorch
 - **Multiple Norms**: Support for 1-norm and 2-norm condition numbers
 - **Rich Algorithm Suite**:
   - **1-norm**: Hager-Higham, Power iteration, Oettli-Prager sampling, Block Hager
-  - **2-norm**: Power method, Lanczos, Arnoldi, Golub-Kahan bidiagonalization
-  - **CuPy integrations**: SVDS, EIGSH, LOBPCG wrappers
+  - **2-norm**: Power method, Lanczos, Golub-Kahan bidiagonalization
+  - **PyTorch integrations**: SVDS, EIGSH, LOBPCG wrappers
 - **Flexible Solver System**: LU, LSMR, CG, GMRES, Direct, Auto-selection
+- **GNN Prediction Module**: Train reusable models that predict condition numbers directly or via inverse-norm prediction
 - **Smart LU Caching**: Reuses factorizations for multiple solves (10-20x speedup)
 - **Memory Efficient**: Designed for large sparse matrices
 
@@ -38,14 +41,14 @@ pip install sparse-kappa
 
 ```bash
 git clone https://github.com/chenxinye/sparse-kappa
-pip install cupy-cuda11x  # or cupy-cuda12x for CUDA 12
+pip install torch
 pip install -e .
 ```
 
 ## Quick Start
 
 ```python
-import cupyx.scipy.sparse as sp
+from sparse_kappa.backend import sparse as sp
 from sparse_kappa import cond_estimate
 
 # Create sparse matrix
@@ -74,16 +77,16 @@ cond = cond_estimate(A, norm=1, method='hager-higham', solver='lu')
 
 ## 2-Norm Methods
 
-| Method       | Description                        | Best For                     | Complexity      |
-|--------------|------------------------------------|------------------------------|-----------------|
-| `svds`       | Partial SVD (most accurate)        | Small-medium matrices (<5k)  | O(k·nnz)        |
-| `eigsh`      | Symmetric eigenvalue solver        | Symmetric matrices           | O(k·nnz)        |
-| `lobpcg`     | Block preconditioned CG            | Large matrices               | O(k·nnz)        |
-| `power`      | Power iteration                    | Quick estimates              | O(k·nnz)        |
-| `lanczos`    | Lanczos tridiagonalization         | Medium matrices              | O(k²·nnz)       |
-| `arnoldi`    | Arnoldi iteration                  | Non-symmetric                | O(k²·nnz)       |
-| `golub-kahan`| Bidiagonalization                  | Numerically stable           | O(k·nnz)        |
-| `auto`       | Automatic selection                | All cases                    | -               |
+| Method            | Description                                                | Best For                                | Complexity      |
+|-------------------|------------------------------------------------------------|-----------------------------------------|-----------------|
+| `svds`            | Partial SVD (most accurate)                                | Small-medium matrices (<5k)             | O(k·nnz)        |
+| `eigsh`           | Symmetric eigenvalue solver                                | Symmetric / Hermitian matrices          | O(k·nnz)        |
+| `lobpcg`          | Block preconditioned CG                                    | Large matrices                          | O(k·nnz)        |
+| `power`           | Power iteration                                            | Quick estimates                         | O(k·nnz)        |
+| `lanczos`         | Lanczos tridiagonalization                                 | Medium symmetric matrices               | O(k²·nnz)       |
+| `lanczos_unsym`   | Lanczos-style condition estimation via `eigsh` on `A^H A`  | Non-symmetric / rectangular matrices    | O(k·nnz)        |
+| `golub-kahan`     | Bidiagonalization                                          | Numerically stable                      | O(k·nnz)        |
+| `auto`            | Automatic selection                                        | All cases                               | -               |
 
 ## Solver Options
 
@@ -102,12 +105,70 @@ All 1-norm methods support flexible solver selection:
 **Legend:**  
 `k` = iterations, `b` = block size, `m` = samples, `nnz` = non-zeros
 
+## GNN-Based Prediction
+
+The `sparse_kappa.gnn` module learns a mapping from sparse matrices to
+condition-number related scalars. It supports two training targets:
+
+- `target="condition"`: predict `kappa(A)` directly.
+- `target="inverse_norm"`: predict `||A^{-1}||`, then compute `||A|| * ||A^{-1}||` at inference time.
+
+The default graph builder turns a sparse matrix into a row/column bipartite
+graph. You can replace the feature extractor, model, optimizer, scheduler,
+loss function, and validation callback.
+
+```python
+from sparse_kappa import TrainingConfig, train_gnn_condition_estimator
+from sparse_kappa.gnn import GNNConditionEstimator
+from sparse_kappa.backend import sparse as sp
+
+train_samples = [
+    {"matrix": A0, "condition_number": 12.3, "norm_Ainv": 0.42},
+    {"matrix": A1, "condition_number": 18.9, "norm_Ainv": 0.61},
+]
+
+# Scheme 1: direct condition-number prediction
+config = TrainingConfig(target="condition", norm=2, epochs=100, lr=1e-3)
+estimator = train_gnn_condition_estimator(
+    train_samples,
+    val_data=None,
+    config=config,
+    save_path="models/gnn_cond.pt",
+)
+
+# Load and predict one matrix or a list of matrices.
+estimator = GNNConditionEstimator.load("models/gnn_cond.pt")
+pred = estimator.predict(sp.random(100, 100, density=0.02, format="csr"))
+
+# Scheme 2: inverse-norm prediction. At inference, sparse-kappa computes ||A||
+# for the requested norm and returns ||A|| * predicted ||A^{-1}||.
+inv_config = TrainingConfig(target="inverse_norm", norm=1, epochs=100)
+inv_estimator = train_gnn_condition_estimator(train_samples, config=inv_config)
+result = inv_estimator.predict(A_test, return_dict=True)
+print(result["condition_number"], result["norm_A"], result["norm_Ainv"])
+```
+
+Customization hooks follow the same shape:
+
+```python
+import torch
+
+estimator.fit(
+    train_samples,
+    val_data=val_samples,
+    optimizer_factory=lambda params: torch.optim.Adam(params, lr=5e-4),
+    scheduler_factory=lambda opt: torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=50),
+    loss_fn=torch.nn.SmoothL1Loss(),
+    validator=my_validation_callback,
+)
+```
+
 ## Examples
 
 ### Example 1: Compare Methods
 
 ```python
-import cupyx.scipy.sparse as sp
+from sparse_kappa.backend import sparse as sp
 from sparse_kappa import cond_estimate
 
 A = sp.random(2000, 2000, density=0.005, format='csr')
@@ -215,10 +276,13 @@ Contributions welcome! Please submit issues and pull requests on GitHub.
 If you use this library in your research, please cite:
 
 ```bibtex
-@software{sparse_kappa,
-  title={Sparse Matrices Condition Number Estimation on GPUs},
-  author={Xinye Chen},
-  year={2026},
-  url={https://github.com/chenxinye/sparse_kappa}
+@misc{carson2026estimatingconditionnumbergraph,
+      title={Estimating condition number with Graph Neural Networks}, 
+      author={Erin Carson and Xinye Chen},
+      year={2026},
+      eprint={2603.10277},
+      archivePrefix={arXiv},
+      primaryClass={cs.LG},
+      url={https://arxiv.org/abs/2603.10277}, 
 }
 ```

{sparse_kappa-0.0.2 → sparse_kappa-0.0.3}/pyproject.toml

@@ -4,8 +4,8 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "sparse-kappa"
-version = "0.0.2"
-description = "GPU-accelerated sparse matrix condition number estimation using CuPy"
+version = "0.0.3"
+description = "PyTorch-accelerated sparse matrix condition number estimation"
 readme = "README.md"
 authors = [
     {name = "Xinye Chen", email = "xinyechenai@gmail.com"}
@@ -24,8 +24,9 @@ classifiers = [
     "Programming Language :: Python :: 3.11",
 ]
 dependencies = [
-    "cupy>=10.0.0",
     "numpy>=1.20.0",
+    "torch>=2.0.0",
+    "myst-parser>=2.0",
 ]
 
 [project.optional-dependencies]
@@ -37,4 +38,4 @@ dev = [
 ]
 
 [project.urls]
-Homepage = "https://github.com/chenxinye/sparse-kappa"
+Homepage = "https://github.com/chenxinye/sparse-kappa"

{sparse_kappa-0.0.2 → sparse_kappa-0.0.3}/setup.py

@@ -5,10 +5,10 @@ with open("README.md", "r", encoding="utf-8") as fh:
 
 setup(
     name="sparse-kappa",                    
-    version="0.0.1",
+    version="0.0.3",
     author="Xinye Chen",
     author_email="xinyechenai@gmail.com",
-    description="GPU-accelerated sparse matrix condition number estimation using CuPy",
+    description="PyTorch-accelerated sparse matrix condition number estimation",
     long_description=long_description,
     long_description_content_type="text/markdown",
     url="https://github.com/chenxinye/sparse-kappa",
@@ -29,8 +29,9 @@ setup(
     ],
     python_requires=">=3.8",
     install_requires=[
-        "cupy>=10.0.0",
+        "torch>=2.0.0",
         "numpy>=1.20.0",
+        "myst-parser>=2.0",
     ],
     extras_require={
         "dev": [

{sparse_kappa-0.0.2 → sparse_kappa-0.0.3}/sparse_kappa/__init__.py

@@ -1,5 +1,5 @@
 """
-CuPy Sparse Condition Number Estimation Library
+PyTorch Sparse Condition Number Estimation Library
 
 A comprehensive GPU-accelerated library for estimating condition numbers
 of sparse matrices using various iterative methods.
@@ -8,14 +8,18 @@ of sparse matrices using various iterative methods.
 from .sparse_kappa import cond_estimate, ConditionNumberEstimator
 from .solvers import create_solver, SOLVER_MAP
 from .config import configure_warnings
+from .gnn import GNNConditionEstimator, TrainingConfig, train_gnn_condition_estimator
 
-__version__ = "0.0.2"
+__version__ = "0.0.3"
 __all__ = [
     "cond_estimate", 
     "ConditionNumberEstimator",
     "create_solver",
-    "SOLVER_MAP"
+    "SOLVER_MAP",
+    "GNNConditionEstimator",
+    "TrainingConfig",
+    "train_gnn_condition_estimator",
 ]
 
 
-configure_warnings(convergence=False, format_conversion=False)
+configure_warnings(convergence=False, format_conversion=False)

sparse_kappa-0.0.3/sparse_kappa/backend/__init__.py

@@ -0,0 +1,6 @@
+"""PyTorch-backed compatibility layer used by sparse-kappa."""
+
+from . import torch_api
+from . import sparse
+
+__all__ = ["torch_api", "sparse"]

sparse_kappa-0.0.3/sparse_kappa/backend/sparse/__init__.py

@@ -0,0 +1,188 @@
+"""Sparse matrix compatibility layer backed by PyTorch tensors."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Optional, Sequence, Tuple
+
+import torch
+
+from .. import torch_api as cp
+from . import linalg
+
+
+@dataclass
+class TorchSparseMatrix:
+    _dense: torch.Tensor
+
+    def __post_init__(self) -> None:
+        if self._dense.ndim != 2:
+            raise ValueError("sparse matrices must be 2D")
+
+    @property
+    def shape(self) -> Tuple[int, int]:
+        return tuple(self._dense.shape)
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return self._dense.dtype
+
+    @property
+    def nnz(self) -> int:
+        return int(torch.count_nonzero(self._dense).item())
+
+    @property
+    def data(self) -> torch.Tensor:
+        return self._dense[self._dense != 0]
+
+    @data.setter
+    def data(self, values: torch.Tensor) -> None:
+        mask = self._dense != 0
+        self._dense[mask] = values.to(dtype=self._dense.dtype, device=self._dense.device)
+
+    @property
+    def T(self) -> "TorchSparseMatrix":
+        return TorchSparseMatrix(self._dense.mT)
+
+    def conj(self) -> "TorchSparseMatrix":
+        return TorchSparseMatrix(torch.conj(self._dense))
+
+    def copy(self) -> "TorchSparseMatrix":
+        return TorchSparseMatrix(self._dense.clone())
+
+    def toarray(self) -> torch.Tensor:
+        return self._dense
+
+    def tocsr(self) -> "TorchSparseMatrix":
+        return self
+
+    def tocsc(self) -> "TorchSparseMatrix":
+        return self
+
+    def astype(self, dtype: Any) -> "TorchSparseMatrix":
+        return TorchSparseMatrix(self._dense.to(dtype=cp._normalize_dtype(dtype)))
+
+    def sum(self, axis: Optional[int] = None):
+        return self._dense.sum() if axis is None else self._dense.sum(dim=axis, keepdim=True)
+
+    def __getitem__(self, key: Any):
+        out = self._dense.__getitem__(key)
+        return TorchSparseMatrix(out) if isinstance(out, torch.Tensor) and out.ndim == 2 else out
+
+    def __matmul__(self, other: Any):
+        rhs = other._dense if isinstance(other, TorchSparseMatrix) else other
+        out = self._dense @ rhs
+        return TorchSparseMatrix(out) if isinstance(out, torch.Tensor) and out.ndim == 2 else out
+
+    def __rmatmul__(self, other: Any):
+        lhs = other._dense if isinstance(other, TorchSparseMatrix) else other
+        out = lhs @ self._dense
+        return TorchSparseMatrix(out) if isinstance(out, torch.Tensor) and out.ndim == 2 else out
+
+    def _binary(self, other: Any, op):
+        rhs = other._dense if isinstance(other, TorchSparseMatrix) else other
+        out = op(self._dense, rhs)
+        return TorchSparseMatrix(out) if isinstance(out, torch.Tensor) and out.ndim == 2 else out
+
+    def __add__(self, other: Any):
+        return self._binary(other, torch.add)
+
+    def __radd__(self, other: Any):
+        return self.__add__(other)
+
+    def __sub__(self, other: Any):
+        return self._binary(other, torch.sub)
+
+    def __rsub__(self, other: Any):
+        lhs = other._dense if isinstance(other, TorchSparseMatrix) else other
+        return TorchSparseMatrix(torch.sub(lhs, self._dense))
+
+    def __mul__(self, other: Any):
+        return self._binary(other, torch.mul)
+
+    def __rmul__(self, other: Any):
+        return self.__mul__(other)
+
+    def __truediv__(self, other: Any):
+        return self._binary(other, torch.true_divide)
+
+    def __repr__(self) -> str:
+        return f"TorchSparseMatrix(shape={self.shape}, nnz={self.nnz}, dtype={self.dtype})"
+
+
+spmatrix = TorchSparseMatrix
+
+
+def _dense_from_any(arg: Any, dtype: Any = None) -> torch.Tensor:
+    if isinstance(arg, TorchSparseMatrix):
+        dense = arg.toarray()
+    elif isinstance(arg, torch.Tensor):
+        dense = arg.to_dense() if arg.is_sparse else arg
+    elif hasattr(arg, "toarray"):
+        dense = cp.asarray(arg.toarray(), dtype=dtype)
+    else:
+        dense = cp.asarray(arg, dtype=dtype)
+    if dtype is not None:
+        dense = dense.to(dtype=cp._normalize_dtype(dtype))
+    return dense
+
+
+def csr_matrix(arg: Any, shape: Optional[Tuple[int, int]] = None, dtype: Any = None) -> TorchSparseMatrix:
+    if isinstance(arg, tuple) and len(arg) == 2:
+        values, coords = arg
+        rows, cols = coords
+        values = cp.asarray(values, dtype=dtype)
+        rows = cp.asarray(rows, dtype=cp.int64)
+        cols = cp.asarray(cols, dtype=cp.int64)
+        if shape is None:
+            shape = (int(rows.max().item()) + 1, int(cols.max().item()) + 1)
+        dense = torch.zeros(shape, dtype=values.dtype, device=values.device)
+        dense[rows.long(), cols.long()] = values
+        return TorchSparseMatrix(dense)
+    dense = _dense_from_any(arg, dtype=dtype)
+    if shape is not None and tuple(dense.shape) != tuple(shape):
+        dense = dense.reshape(shape)
+    return TorchSparseMatrix(dense.clone())
+
+
+def coo_matrix(arg: Any, shape: Optional[Tuple[int, int]] = None, dtype: Any = None) -> TorchSparseMatrix:
+    return csr_matrix(arg, shape=shape, dtype=dtype)
+
+
+def isspmatrix(x: Any) -> bool:
+    return isinstance(x, TorchSparseMatrix)
+
+
+def isspmatrix_csr(x: Any) -> bool:
+    return isinstance(x, TorchSparseMatrix)
+
+
+def eye(n: int, m: Optional[int] = None, dtype: Any = cp.float64, format: str = "csr") -> TorchSparseMatrix:
+    return TorchSparseMatrix(cp.eye(n, m, dtype=dtype))
+
+
+def identity(n: int, dtype: Any = cp.float64, format: str = "csr") -> TorchSparseMatrix:
+    return eye(n, dtype=dtype, format=format)
+
+
+def diags(diagonals: Any, offsets: int = 0, shape: Optional[Tuple[int, int]] = None, format: str = "csr", dtype: Any = None) -> TorchSparseMatrix:
+    diag = cp.asarray(diagonals, dtype=dtype)
+    if offsets != 0:
+        raise NotImplementedError("Only the main diagonal is supported")
+    n = diag.numel()
+    if shape is None:
+        shape = (n, n)
+    dense = torch.zeros(shape, dtype=diag.dtype, device=diag.device)
+    idx = torch.arange(n, device=diag.device)
+    dense[idx, idx] = diag
+    return TorchSparseMatrix(dense)
+
+
+def random(m: int, n: int, density: float = 0.01, format: str = "csr", dtype: Any = cp.float64, random_state: Optional[int] = None) -> TorchSparseMatrix:
+    dtype = cp._normalize_dtype(dtype)
+    generator = torch.Generator(device=cp._default_device())
+    if random_state is not None:
+        generator.manual_seed(random_state)
+    mask = torch.rand((m, n), generator=generator, device=cp._default_device()) < density
+    values = torch.rand((m, n), generator=generator, device=cp._default_device(), dtype=dtype)
+    return TorchSparseMatrix(torch.where(mask, values, torch.zeros((), dtype=dtype, device=values.device)))