torchsparsegradutils 0.2.0__tar.gz → 0.2.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (81) hide show
  1. {torchsparsegradutils-0.2.0/torchsparsegradutils.egg-info → torchsparsegradutils-0.2.2}/PKG-INFO +38 -31
  2. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/README.md +14 -10
  3. torchsparsegradutils-0.2.2/docs/source/conf.py +187 -0
  4. torchsparsegradutils-0.2.2/pyproject.toml +116 -0
  5. torchsparsegradutils-0.2.2/setup.py +5 -0
  6. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/cupy/__init__.py +11 -2
  7. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/cupy/cupy_bindings.py +25 -11
  8. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/cupy/cupy_sparse_solve.py +92 -95
  9. torchsparsegradutils-0.2.2/torchsparsegradutils/encoders/__init__.py +44 -0
  10. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/encoders/pairwise_encoder.py +6 -2
  11. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/encoders/pairwise_voxel_encoder.py +3 -10
  12. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/indexed_matmul.py +3 -2
  13. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/sparse_matmul.py +3 -3
  14. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/sparse_solve.py +2 -2
  15. torchsparsegradutils-0.2.2/torchsparsegradutils/tests/conftest.py +48 -0
  16. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/tests/test_bicgstab.py +8 -11
  17. torchsparsegradutils-0.2.2/torchsparsegradutils/tests/test_config.py +62 -0
  18. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/tests/test_cupy_bindings.py +9 -12
  19. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/tests/test_cupy_sparse_solve.py +45 -39
  20. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/tests/test_dist_stats_helpers.py +1 -12
  21. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/tests/test_distributions.py +22 -42
  22. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/tests/test_encoders.py +52 -60
  23. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/tests/test_indexed_matmul.py +30 -5
  24. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/tests/test_integration_pairwise_sparse_mvn.py +68 -18
  25. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/tests/test_jax_bindings.py +1 -5
  26. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/tests/test_jax_sparse_solve.py +19 -27
  27. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/tests/test_linear_cg.py +0 -19
  28. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/tests/test_lsmr.py +19 -23
  29. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/tests/test_minres.py +3 -11
  30. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/tests/test_random.py +9 -10
  31. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/tests/test_sparse_lstsq.py +3 -7
  32. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/tests/test_sparse_matmul.py +7 -14
  33. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/tests/test_sparse_solve.py +28 -37
  34. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/tests/test_sparse_triangular_solve.py +8 -17
  35. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/tests/test_utils.py +2 -7
  36. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/utils/utils.py +6 -8
  37. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2/torchsparsegradutils.egg-info}/PKG-INFO +38 -31
  38. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils.egg-info/SOURCES.txt +3 -0
  39. torchsparsegradutils-0.2.2/torchsparsegradutils.egg-info/requires.txt +31 -0
  40. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils.egg-info/top_level.txt +3 -0
  41. torchsparsegradutils-0.2.0/pyproject.toml +0 -49
  42. torchsparsegradutils-0.2.0/setup.py +0 -50
  43. torchsparsegradutils-0.2.0/torchsparsegradutils/encoders/__init__.py +0 -14
  44. torchsparsegradutils-0.2.0/torchsparsegradutils.egg-info/requires.txt +0 -15
  45. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/LICENSE +0 -0
  46. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/MANIFEST.in +0 -0
  47. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/setup.cfg +0 -0
  48. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/__init__.py +0 -0
  49. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/benchmarks/__init__.py +0 -0
  50. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/benchmarks/batched_sparse_mm_rand.py +0 -0
  51. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/benchmarks/benchmark_suite.py +0 -0
  52. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/benchmarks/benchmark_utils.py +0 -0
  53. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/benchmarks/sparse_generic_solve_rand.py +0 -0
  54. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/benchmarks/sparse_generic_solve_suite.py +0 -0
  55. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/benchmarks/sparse_mm_rand.py +0 -0
  56. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/benchmarks/sparse_mm_suite.py +0 -0
  57. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/benchmarks/sparse_triangular_solve_rand.py +0 -0
  58. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/benchmarks/sparse_triangular_solve_suitesparse.py +0 -0
  59. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/benchmarks/visualize_benchmark_results.py +0 -0
  60. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/distributions/__init__.py +0 -0
  61. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/distributions/constraints.py +0 -0
  62. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/distributions/sparse_multivariate_normal.py +0 -0
  63. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/jax/__init__.py +0 -0
  64. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/jax/jax_bindings.py +0 -0
  65. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/jax/jax_sparse_solve.py +0 -0
  66. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/sparse_lstsq.py +0 -0
  67. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/tests/__init__.py +0 -0
  68. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/tests/test_doctests.py +0 -0
  69. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/tests/test_params/czyx_shifts.yaml +0 -0
  70. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/tests/test_params/pairwise_coo_indices.yaml +0 -0
  71. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/tests/test_params/xyz_coords.yaml +0 -0
  72. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/tests/test_quickstart_guide.py +0 -0
  73. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/utils/__init__.py +0 -0
  74. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/utils/bicgstab.py +0 -0
  75. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/utils/dist_stats_helpers.py +0 -0
  76. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/utils/linear_cg.py +0 -0
  77. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/utils/lsmr.py +0 -0
  78. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/utils/minres.py +0 -0
  79. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils/utils/random_sparse.py +0 -0
  80. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils.egg-info/dependency_links.txt +0 -0
  81. {torchsparsegradutils-0.2.0 → torchsparsegradutils-0.2.2}/torchsparsegradutils.egg-info/not-zip-safe +0 -0
@@ -1,14 +1,16 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: torchsparsegradutils
3
- Version: 0.2.0
3
+ Version: 0.2.2
4
4
  Summary: A collection of utility functions to work with PyTorch sparse tensors
5
- Home-page: https://github.com/cai4cai/torchsparsegradutils
6
- Author: CAI4CAI research group
7
- Author-email: contact@cai4cai.uk
8
- License: Apache-2.0
9
- Keywords: sparse torch utility
5
+ Author-email: CAI4CAI research group <contact@cai4cai.uk>
6
+ License-Expression: Apache-2.0
7
+ Project-URL: Homepage, https://github.com/cai4cai/torchsparsegradutils
8
+ Project-URL: Documentation, https://torchsparsegradutils.readthedocs.io
9
+ Project-URL: Repository, https://github.com/cai4cai/torchsparsegradutils
10
+ Project-URL: Issues, https://github.com/cai4cai/torchsparsegradutils/issues
11
+ Project-URL: Changelog, https://github.com/cai4cai/torchsparsegradutils/releases
12
+ Keywords: sparse,torch,utility
10
13
  Classifier: Operating System :: OS Independent
11
- Classifier: License :: OSI Approved :: Apache Software License
12
14
  Classifier: Programming Language :: Python :: 3.10
13
15
  Classifier: Programming Language :: Python :: 3.11
14
16
  Classifier: Programming Language :: Python :: 3.12
@@ -17,9 +19,14 @@ Description-Content-Type: text/markdown
17
19
  License-File: LICENSE
18
20
  Requires-Dist: torch>=2.5
19
21
  Requires-Dist: scipy
20
- Provides-Extra: extras
21
- Requires-Dist: jax; extra == "extras"
22
- Requires-Dist: cupy; extra == "extras"
22
+ Requires-Dist: packaging
23
+ Provides-Extra: cupy
24
+ Requires-Dist: cupy-cuda12x>=13.0; extra == "cupy"
25
+ Provides-Extra: jax
26
+ Requires-Dist: jax[cuda12]; extra == "jax"
27
+ Provides-Extra: all
28
+ Requires-Dist: cupy-cuda12x>=13.0; extra == "all"
29
+ Requires-Dist: jax[cuda12]; extra == "all"
23
30
  Provides-Extra: docs
24
31
  Requires-Dist: sphinx>=7.0.0; extra == "docs"
25
32
  Requires-Dist: sphinx-rtd-theme>=1.3.0; extra == "docs"
@@ -28,23 +35,19 @@ Requires-Dist: myst-parser>=2.0.0; extra == "docs"
28
35
  Requires-Dist: sphinx-autobuild>=2021.3.14; extra == "docs"
29
36
  Requires-Dist: matplotlib>=3.5.0; extra == "docs"
30
37
  Requires-Dist: sphinx-autodoc-typehints>=1.24.0; extra == "docs"
31
- Dynamic: author
32
- Dynamic: author-email
33
- Dynamic: classifier
34
- Dynamic: description
35
- Dynamic: description-content-type
36
- Dynamic: home-page
37
- Dynamic: keywords
38
- Dynamic: license
38
+ Requires-Dist: sphinxcontrib-bibtex>=2.5.0; extra == "docs"
39
+ Requires-Dist: sphinx-gallery>=0.13.0; extra == "docs"
40
+ Provides-Extra: dev
41
+ Requires-Dist: pytest; extra == "dev"
42
+ Requires-Dist: pytest-cov; extra == "dev"
43
+ Requires-Dist: black; extra == "dev"
44
+ Requires-Dist: isort; extra == "dev"
45
+ Requires-Dist: flake8; extra == "dev"
39
46
  Dynamic: license-file
40
- Dynamic: provides-extra
41
- Dynamic: requires-dist
42
- Dynamic: requires-python
43
- Dynamic: summary
44
47
 
45
48
  # torchsparsegradutils: Sparsity-preserving gradient utility tools for PyTorch
46
49
 
47
- [![Python tests](https://github.com/cai4cai/torchsparsegradutils/actions/workflows/python-package.yml/badge.svg)](https://github.com/cai4cai/torchsparsegradutils/actions/workflows/python-package.yml) [![License](https://img.shields.io/github/license/cai4cai/torchsparsegradutils)](https://github.com/cai4cai/torchsparsegradutils?tab=Apache-2.0-1-ov-file#readme) [![Code Style: Black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
50
+ [![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)
48
51
 
49
52
  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.
50
53
 
@@ -135,17 +138,21 @@ pip install git+https://github.com/cai4cai/torchsparsegradutils
135
138
  For full functionality, install optional dependencies:
136
139
 
137
140
  ```bash
138
- # For CuPy sparse solver support (GPU acceleration)
139
- pip install cupy-cuda12x # Replace with your CUDA version
141
+ # For CuPy sparse solver support (GPU acceleration, requires CUDA 12.x)
142
+ pip install torchsparsegradutils[cupy]
140
143
 
141
144
  # For JAX sparse solver support
142
- pip install "jax[cpu]" # CPU version
143
- pip install "jax[cuda12]" # GPU version (replace with your CUDA version)
145
+ pip install torchsparsegradutils[jax]
146
+
147
+ # Install all optional dependencies
148
+ pip install torchsparsegradutils[all]
144
149
 
145
150
  # For benchmarking and testing
146
151
  pip install scipy matplotlib pandas tqdm pytest
147
152
  ```
148
153
 
154
+ > **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`).
155
+
149
156
  ### Requirements
150
157
 
151
158
  - **Python**: ≥ 3.10
@@ -156,13 +163,13 @@ pip install scipy matplotlib pandas tqdm pytest
156
163
 
157
164
  ## 📊 Performance Benchmarks
158
165
 
159
- Our comprehensive benchmark suite demonstrates significant performance improvements across various sparse operations. All benchmarks were conducted on an NVIDIA GeForce RTX 4090 with PyTorch 2.8.0+cu128. Benchmarks are performed using [Rothberg/cfd2](https://suitesparse-collection-website.herokuapp.com/Rothberg/cfd2) matrix from [SuiteSparse Matrix Collection](https://suitesparse-collection-website.herokuapp.com/)
166
+ Our comprehensive benchmark suite demonstrates significant performance improvements across various sparse operations. All benchmarks were conducted on an NVIDIA GeForce RTX 4090 with PyTorch 2.9.0+cu130. Benchmarks are performed using [Rothberg/cfd2](https://suitesparse-collection-website.herokuapp.com/Rothberg/cfd2) matrix from [SuiteSparse Matrix Collection](https://suitesparse-collection-website.herokuapp.com/)
160
167
 
161
168
  ![Sparse MM Suite Performance (int32/float32 COO)](torchsparsegradutils/benchmarks/benchmark_visualizations/sparse_mm_suite_performance_int32_float32_coo.png)
162
169
 
163
170
  ![Sparse Triangular Solve Suite Performance (int32/float32 COO)](torchsparsegradutils/benchmarks/benchmark_visualizations/triangular_solve_suitesparse_performance_int32_float32_coo.png)
164
171
 
165
- ![Sparse Genertic Solve Suite Performance (int32/float32 COO)](torchsparsegradutils/benchmarks/benchmark_visualizations/sparse_solve_suite_performance_int32_float32_coo.png)
172
+ ![Sparse Generic Solve Suite Performance (int32/float32 COO)](torchsparsegradutils/benchmarks/benchmark_visualizations/sparse_solve_suite_performance_int32_float32_coo.png)
166
173
 
167
174
  ## 🚀 Quick Start
168
175
 
@@ -543,7 +550,7 @@ For a consistent development environment with GPU support and all dependencies p
543
550
 
544
551
  **Available Configurations:**
545
552
 
546
- - **`.devcontainer/Dockerfile.stable`** (default): Uses stable PyTorch with CUDA 12.8 support
553
+ - **`.devcontainer/Dockerfile.stable`** (default): Uses stable PyTorch with CUDA 13.0 support
547
554
  - **`.devcontainer/Dockerfile.nightly`**: Uses nightly PyTorch builds for latest features
548
555
 
549
556
  To switch configurations, modify the `dockerfile` field in `.devcontainer/devcontainer.json`:
@@ -555,7 +562,7 @@ To switch configurations, modify the `dockerfile` field in `.devcontainer/devcon
555
562
  ```
556
563
 
557
564
  **What's Included:**
558
- - **CUDA 12.8**: Full GPU development support with NVIDIA drivers
565
+ - **CUDA 13.0**: Full GPU development support with NVIDIA drivers
559
566
  - **Pre-installed Dependencies**: PyTorch, CuPy, JAX, SciPy, and all development tools
560
567
  - **VS Code Extensions**: Python, Pylance, Jupyter, GitHub Copilot, and code formatting tools
561
568
  - **Development Tools**: pytest, black, flake8, pre-commit hooks
@@ -1,6 +1,6 @@
1
1
  # torchsparsegradutils: Sparsity-preserving gradient utility tools for PyTorch
2
2
 
3
- [![Python tests](https://github.com/cai4cai/torchsparsegradutils/actions/workflows/python-package.yml/badge.svg)](https://github.com/cai4cai/torchsparsegradutils/actions/workflows/python-package.yml) [![License](https://img.shields.io/github/license/cai4cai/torchsparsegradutils)](https://github.com/cai4cai/torchsparsegradutils?tab=Apache-2.0-1-ov-file#readme) [![Code Style: Black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
3
+ [![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)
4
4
 
5
5
  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.
6
6
 
@@ -91,17 +91,21 @@ pip install git+https://github.com/cai4cai/torchsparsegradutils
91
91
  For full functionality, install optional dependencies:
92
92
 
93
93
  ```bash
94
- # For CuPy sparse solver support (GPU acceleration)
95
- pip install cupy-cuda12x # Replace with your CUDA version
94
+ # For CuPy sparse solver support (GPU acceleration, requires CUDA 12.x)
95
+ pip install torchsparsegradutils[cupy]
96
96
 
97
97
  # For JAX sparse solver support
98
- pip install "jax[cpu]" # CPU version
99
- pip install "jax[cuda12]" # GPU version (replace with your CUDA version)
98
+ pip install torchsparsegradutils[jax]
99
+
100
+ # Install all optional dependencies
101
+ pip install torchsparsegradutils[all]
100
102
 
101
103
  # For benchmarking and testing
102
104
  pip install scipy matplotlib pandas tqdm pytest
103
105
  ```
104
106
 
107
+ > **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`).
108
+
105
109
  ### Requirements
106
110
 
107
111
  - **Python**: ≥ 3.10
@@ -112,13 +116,13 @@ pip install scipy matplotlib pandas tqdm pytest
112
116
 
113
117
  ## 📊 Performance Benchmarks
114
118
 
115
- Our comprehensive benchmark suite demonstrates significant performance improvements across various sparse operations. All benchmarks were conducted on an NVIDIA GeForce RTX 4090 with PyTorch 2.8.0+cu128. Benchmarks are performed using [Rothberg/cfd2](https://suitesparse-collection-website.herokuapp.com/Rothberg/cfd2) matrix from [SuiteSparse Matrix Collection](https://suitesparse-collection-website.herokuapp.com/)
119
+ Our comprehensive benchmark suite demonstrates significant performance improvements across various sparse operations. All benchmarks were conducted on an NVIDIA GeForce RTX 4090 with PyTorch 2.9.0+cu130. Benchmarks are performed using [Rothberg/cfd2](https://suitesparse-collection-website.herokuapp.com/Rothberg/cfd2) matrix from [SuiteSparse Matrix Collection](https://suitesparse-collection-website.herokuapp.com/)
116
120
 
117
121
  ![Sparse MM Suite Performance (int32/float32 COO)](torchsparsegradutils/benchmarks/benchmark_visualizations/sparse_mm_suite_performance_int32_float32_coo.png)
118
122
 
119
123
  ![Sparse Triangular Solve Suite Performance (int32/float32 COO)](torchsparsegradutils/benchmarks/benchmark_visualizations/triangular_solve_suitesparse_performance_int32_float32_coo.png)
120
124
 
121
- ![Sparse Genertic Solve Suite Performance (int32/float32 COO)](torchsparsegradutils/benchmarks/benchmark_visualizations/sparse_solve_suite_performance_int32_float32_coo.png)
125
+ ![Sparse Generic Solve Suite Performance (int32/float32 COO)](torchsparsegradutils/benchmarks/benchmark_visualizations/sparse_solve_suite_performance_int32_float32_coo.png)
122
126
 
123
127
  ## 🚀 Quick Start
124
128
 
@@ -499,7 +503,7 @@ For a consistent development environment with GPU support and all dependencies p
499
503
 
500
504
  **Available Configurations:**
501
505
 
502
- - **`.devcontainer/Dockerfile.stable`** (default): Uses stable PyTorch with CUDA 12.8 support
506
+ - **`.devcontainer/Dockerfile.stable`** (default): Uses stable PyTorch with CUDA 13.0 support
503
507
  - **`.devcontainer/Dockerfile.nightly`**: Uses nightly PyTorch builds for latest features
504
508
 
505
509
  To switch configurations, modify the `dockerfile` field in `.devcontainer/devcontainer.json`:
@@ -511,7 +515,7 @@ To switch configurations, modify the `dockerfile` field in `.devcontainer/devcon
511
515
  ```
512
516
 
513
517
  **What's Included:**
514
- - **CUDA 12.8**: Full GPU development support with NVIDIA drivers
518
+ - **CUDA 13.0**: Full GPU development support with NVIDIA drivers
515
519
  - **Pre-installed Dependencies**: PyTorch, CuPy, JAX, SciPy, and all development tools
516
520
  - **VS Code Extensions**: Python, Pylance, Jupyter, GitHub Copilot, and code formatting tools
517
521
  - **Development Tools**: pytest, black, flake8, pre-commit hooks
@@ -658,4 +662,4 @@ dist_stable = SparseMultivariateNormal(
658
662
  - **No SPD Constraints**: Doesn't require strict positive definiteness
659
663
  - **Better Conditioning**: Diagonal component can be controlled independently
660
664
 
661
- **Status**: This is a known limitation of the LL^T precision formulation. LDL^T parameterization is the recommended approach for precision matrices.
665
+ **Status**: This is a known limitation of the LL^T precision formulation. LDL^T parameterization is the recommended approach for precision matrices.
@@ -0,0 +1,187 @@
1
+ # Configuration file for the Sphinx documentation builder.
2
+ #
3
+ # For the full list of built-in configuration values, see the documentation:
4
+ # https://www.sphinx-doc.org/en/master/usage/configuration.html
5
+
6
+ # -- Project information -----------------------------------------------------
7
+ # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
8
+
9
+ import os
10
+ import sys
11
+
12
+ # Add the package to the Python path
13
+ sys.path.insert(0, os.path.abspath("../../"))
14
+
15
+ project = "torchsparsegradutils"
16
+ copyright = "2026, CAI4CAI research group"
17
+ author = "CAI4CAI research group"
18
+ release = "0.2.2"
19
+ version = "0.2.2"
20
+
21
+ # -- General configuration ---------------------------------------------------
22
+ # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration
23
+
24
+ extensions = [
25
+ "sphinx.ext.autodoc", # Automatic documentation from docstrings
26
+ "sphinx.ext.autosummary", # Generate autodoc summaries
27
+ "sphinx.ext.napoleon", # Support for NumPy and Google style docstrings
28
+ "sphinx.ext.viewcode", # Add links to highlighted source code
29
+ "sphinx.ext.mathjax", # Math notation support
30
+ "sphinx.ext.intersphinx", # Link to other project's documentation
31
+ "sphinx.ext.coverage", # Documentation coverage
32
+ "sphinx.ext.githubpages", # GitHub Pages support
33
+ "sphinx_copybutton", # Add copy button to code blocks
34
+ "myst_parser", # Markdown support
35
+ ]
36
+
37
+ # Autosummary settings
38
+ autosummary_generate = True
39
+ autosummary_imported_members = True
40
+
41
+ # Napoleon settings for docstring parsing
42
+ napoleon_google_docstring = True
43
+ napoleon_numpy_docstring = True
44
+ napoleon_include_init_with_doc = False
45
+ napoleon_include_private_with_doc = False
46
+ napoleon_include_special_with_doc = True
47
+ napoleon_use_admonition_for_examples = False
48
+ napoleon_use_admonition_for_notes = False
49
+ napoleon_use_admonition_for_references = False
50
+ napoleon_use_ivar = False
51
+ napoleon_use_param = True
52
+ napoleon_use_rtype = True
53
+ napoleon_preprocess_types = False
54
+ napoleon_type_aliases = None
55
+ napoleon_attr_annotations = True
56
+
57
+ # Autodoc settings
58
+ autodoc_default_options = {
59
+ "members": True,
60
+ "member-order": "bysource",
61
+ "special-members": "__init__",
62
+ "undoc-members": True,
63
+ "exclude-members": "__weakref__",
64
+ }
65
+
66
+ # Intersphinx mapping
67
+ intersphinx_mapping = {
68
+ "python": ("https://docs.python.org/3/", None),
69
+ "torch": ("https://pytorch.org/docs/stable/", None),
70
+ "numpy": ("https://numpy.org/doc/stable/", None),
71
+ "scipy": ("https://docs.scipy.org/doc/scipy/", None),
72
+ "jax": ("https://jax.readthedocs.io/en/latest/", None),
73
+ "cupy": ("https://docs.cupy.dev/en/stable/", None),
74
+ }
75
+
76
+ templates_path = ["_templates"]
77
+ exclude_patterns = []
78
+
79
+ # Source file parsers
80
+ source_suffix = {
81
+ ".rst": None,
82
+ ".md": "myst_parser",
83
+ }
84
+
85
+ # -- Options for HTML output -------------------------------------------------
86
+ # https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output
87
+
88
+ html_theme = "sphinx_rtd_theme"
89
+ html_static_path = ["_static"]
90
+
91
+ # Theme options
92
+ html_theme_options = {
93
+ "canonical_url": "",
94
+ "analytics_id": "",
95
+ "logo_only": False,
96
+ "prev_next_buttons_location": "bottom",
97
+ "style_external_links": False,
98
+ "vcs_pageview_mode": "",
99
+ "style_nav_header_background": "#2980B9",
100
+ # Toc options
101
+ "collapse_navigation": True,
102
+ "sticky_navigation": True,
103
+ "navigation_depth": 4,
104
+ "includehidden": True,
105
+ "titles_only": False,
106
+ }
107
+
108
+ # Add any paths that contain custom static files (such as style sheets) here,
109
+ # relative to this directory. They are copied after the builtin static files,
110
+ # so a file named "default.css" will overwrite the builtin "default.css".
111
+ html_static_path = ["_static"]
112
+
113
+ # Custom CSS
114
+ html_css_files = [
115
+ "custom.css",
116
+ ]
117
+
118
+ # HTML title
119
+ html_title = f"{project} v{version}"
120
+
121
+ # Favicon
122
+ # html_favicon = '_static/favicon.ico'
123
+
124
+ # Logo
125
+ # html_logo = '_static/logo.png'
126
+
127
+ # -- Options for LaTeX output ------------------------------------------------
128
+ latex_engine = "pdflatex"
129
+ latex_elements = {
130
+ # The paper size ('letterpaper' or 'a4paper').
131
+ "papersize": "letterpaper",
132
+ # The font size ('10pt', '11pt' or '12pt').
133
+ "pointsize": "10pt",
134
+ # Additional stuff for the LaTeX preamble.
135
+ "preamble": r"""
136
+ \usepackage{amsmath}
137
+ \usepackage{amsfonts}
138
+ \usepackage{amssymb}
139
+ \usepackage{mathrsfs}
140
+ """,
141
+ }
142
+
143
+ # Grouping the document tree into LaTeX files. List of tuples
144
+ # (source start file, target name, title, author, documentclass [howto, manual, or own class]).
145
+ latex_documents = [
146
+ ("index", "torchsparsegradutils.tex", "torchsparsegradutils Documentation", "CAI4CAI research group", "manual"),
147
+ ]
148
+
149
+ # -- Options for manual page output ------------------------------------------
150
+ # One entry per manual page. List of tuples
151
+ # (source start file, name, description, authors, manual section).
152
+ man_pages = [("index", "torchsparsegradutils", "torchsparsegradutils Documentation", [author], 1)]
153
+
154
+ # -- Options for Texinfo output ----------------------------------------------
155
+ # Grouping the document tree into Texinfo files. List of tuples
156
+ # (source start file, target name, title, author,
157
+ # dir menu entry, description, category)
158
+ texinfo_documents = [
159
+ (
160
+ "index",
161
+ "torchsparsegradutils",
162
+ "torchsparsegradutils Documentation",
163
+ author,
164
+ "torchsparsegradutils",
165
+ "PyTorch sparse tensor utilities.",
166
+ "Miscellaneous",
167
+ ),
168
+ ]
169
+
170
+ # -- Extension configuration -------------------------------------------------
171
+ # Copy button configuration
172
+ copybutton_prompt_text = r">>> |\.\.\. |\$ |In \[\d*\]: | {2,5}\.\.\.: | {5,8}: "
173
+ copybutton_prompt_is_regexp = True
174
+
175
+ # Math notation
176
+ mathjax3_config = {
177
+ "tex": {
178
+ "inlineMath": [["$", "$"], ["\\(", "\\)"]],
179
+ "displayMath": [["$$", "$$"], ["\\[", "\\]"]],
180
+ "processEscapes": True,
181
+ "processEnvironments": True,
182
+ },
183
+ "options": {
184
+ "ignoreHtmlClass": "tex2jax_ignore",
185
+ "processHtmlClass": "tex2jax_process",
186
+ },
187
+ }
@@ -0,0 +1,116 @@
1
+ [build-system]
2
+ requires = ["setuptools>=61.0", "wheel"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "torchsparsegradutils"
7
+ version = "0.2.2"
8
+ description = "A collection of utility functions to work with PyTorch sparse tensors"
9
+ readme = "README.md"
10
+ requires-python = ">=3.10"
11
+ license = "Apache-2.0"
12
+ authors = [
13
+ {name = "CAI4CAI research group", email = "contact@cai4cai.uk"}
14
+ ]
15
+ keywords = ["sparse", "torch", "utility"]
16
+ classifiers = [
17
+ "Operating System :: OS Independent",
18
+ "Programming Language :: Python :: 3.10",
19
+ "Programming Language :: Python :: 3.11",
20
+ "Programming Language :: Python :: 3.12",
21
+ ]
22
+ dependencies = [
23
+ "torch>=2.5",
24
+ "scipy",
25
+ "packaging",
26
+ ]
27
+
28
+ [project.optional-dependencies]
29
+ cupy = ["cupy-cuda12x>=13.0"]
30
+ jax = ["jax[cuda12]"]
31
+ all = ["cupy-cuda12x>=13.0", "jax[cuda12]"]
32
+ docs = [
33
+ "sphinx>=7.0.0",
34
+ "sphinx-rtd-theme>=1.3.0",
35
+ "sphinx-copybutton>=0.5.0",
36
+ "myst-parser>=2.0.0",
37
+ "sphinx-autobuild>=2021.3.14",
38
+ "matplotlib>=3.5.0",
39
+ "sphinx-autodoc-typehints>=1.24.0",
40
+ "sphinxcontrib-bibtex>=2.5.0",
41
+ "sphinx-gallery>=0.13.0",
42
+ ]
43
+ dev = [
44
+ "pytest",
45
+ "pytest-cov",
46
+ "black",
47
+ "isort",
48
+ "flake8",
49
+ ]
50
+
51
+ [project.urls]
52
+ Homepage = "https://github.com/cai4cai/torchsparsegradutils"
53
+ Documentation = "https://torchsparsegradutils.readthedocs.io"
54
+ Repository = "https://github.com/cai4cai/torchsparsegradutils"
55
+ Issues = "https://github.com/cai4cai/torchsparsegradutils/issues"
56
+ Changelog = "https://github.com/cai4cai/torchsparsegradutils/releases"
57
+
58
+ [tool.setuptools]
59
+ zip-safe = false
60
+ include-package-data = true
61
+
62
+ [tool.setuptools.packages.find]
63
+ exclude = ["tests*"]
64
+
65
+ [tool.black]
66
+ line-length = 120
67
+ target-version = ['py310', 'py311', 'py312']
68
+ include = '\.pyi?$'
69
+ exclude = '''
70
+ (
71
+ /(
72
+ # exclude a few common directories in the root of the project
73
+ \.eggs
74
+ | \.git
75
+ | \.hg
76
+ | \.mypy_cache
77
+ | \.tox
78
+ | \.venv
79
+ | \.pytype
80
+ | _build
81
+ | buck-out
82
+ | build
83
+ | dist
84
+ )/
85
+ # also separately exclude other files if needed
86
+ #| some_file
87
+ )
88
+ '''
89
+
90
+ [tool.isort]
91
+ profile = "black"
92
+ multi_line_output = 3
93
+ line_length = 120
94
+ known_first_party = ["torchsparsegradutils"]
95
+ force_grid_wrap = 0
96
+ combine_as_imports = true
97
+ include_trailing_comma = true
98
+
99
+ [tool.pytest.ini_options]
100
+ minversion = "6.0"
101
+ addopts = ["-v"]
102
+ testpaths = [
103
+ "torchsparsegradutils/tests",
104
+ ]
105
+ pythonpath = [
106
+ "torchsparsegradutils/tests",
107
+ ]
108
+ doctest_optionflags = [
109
+ "NORMALIZE_WHITESPACE",
110
+ "IGNORE_EXCEPTION_DETAIL",
111
+ "ELLIPSIS"
112
+ ]
113
+ # Separate command for doctests due to potential issues
114
+ markers = [
115
+ "slow: marks tests as slow (deselect with '-m \"not slow\"')",
116
+ ]
@@ -0,0 +1,5 @@
1
+ # Minimal setup.py for backward compatibility
2
+ # All configuration is now in pyproject.toml
3
+ from setuptools import setup
4
+
5
+ setup()
@@ -11,7 +11,16 @@ if cupy_spec is None:
11
11
  else:
12
12
  have_cupy = True
13
13
 
14
- from .cupy_bindings import _get_array_modules, c2t_coo, c2t_csr, t2c_coo, t2c_csr
14
+ from .cupy_bindings import _backend_to_torch, _get_array_modules, _torch_to_backend, c2t_coo, c2t_csr, t2c_coo, t2c_csr
15
15
  from .cupy_sparse_solve import sparse_solve_c4t
16
16
 
17
- __all__ = ["c2t_coo", "t2c_coo", "c2t_csr", "t2c_csr", "_get_array_modules", "sparse_solve_c4t"]
17
+ __all__ = [
18
+ "c2t_coo",
19
+ "t2c_coo",
20
+ "c2t_csr",
21
+ "t2c_csr",
22
+ "_get_array_modules",
23
+ "_torch_to_backend",
24
+ "_backend_to_torch",
25
+ "sparse_solve_c4t",
26
+ ]
@@ -38,6 +38,20 @@ import scipy.sparse.linalg
38
38
  import torch
39
39
 
40
40
 
41
+ def _torch_to_backend(tensor: torch.Tensor, xp) -> Any:
42
+ r"""Convert a PyTorch tensor to a backend array (NumPy or CuPy) using DLPack for CUDA."""
43
+ if tsgucupy.have_cupy and xp is cp:
44
+ return cp.from_dlpack(tensor)
45
+ return xp.asarray(tensor)
46
+
47
+
48
+ def _backend_to_torch(array) -> torch.Tensor:
49
+ r"""Convert a backend array (NumPy or CuPy) to a PyTorch tensor using DLPack for CUDA."""
50
+ if tsgucupy.have_cupy and isinstance(array, cp.ndarray):
51
+ return torch.from_dlpack(array)
52
+ return torch.as_tensor(array)
53
+
54
+
41
55
  def _get_array_modules(x: Any) -> Tuple[Any, Any]:
42
56
  r"""
43
57
  Select dense & sparse array modules (NumPy/SciPy or CuPy/cupyx.scipy.sparse).
@@ -116,9 +130,9 @@ def t2c_csr(x_torch: torch.Tensor) -> Any:
116
130
  """
117
131
  xp, xsp = _get_array_modules(x_torch)
118
132
 
119
- data_c = xp.asarray(x_torch.values())
120
- col_idx_c = xp.asarray(x_torch.col_indices())
121
- ind_ptr_c = xp.asarray(x_torch.crow_indices())
133
+ data_c = _torch_to_backend(x_torch.values(), xp)
134
+ col_idx_c = _torch_to_backend(x_torch.col_indices(), xp)
135
+ ind_ptr_c = _torch_to_backend(x_torch.crow_indices(), xp)
122
136
  x_cupy = xsp.csr_matrix((data_c, col_idx_c, ind_ptr_c), shape=x_torch.shape)
123
137
  return x_cupy
124
138
 
@@ -150,9 +164,9 @@ def c2t_csr(x_cupy: Any) -> torch.Tensor:
150
164
  >>> x.layout is torch.sparse_csr
151
165
  True
152
166
  """
153
- data_t = torch.as_tensor(x_cupy.data)
154
- idices_t = torch.as_tensor(x_cupy.indices)
155
- ind_ptr_t = torch.as_tensor(x_cupy.indptr)
167
+ data_t = _backend_to_torch(x_cupy.data)
168
+ idices_t = _backend_to_torch(x_cupy.indices)
169
+ ind_ptr_t = _backend_to_torch(x_cupy.indptr)
156
170
  x_torch = torch.sparse_csr_tensor(ind_ptr_t, idices_t, data_t, x_cupy.shape)
157
171
  return x_torch
158
172
 
@@ -205,8 +219,8 @@ def t2c_coo(x_torch: torch.Tensor) -> Any:
205
219
  "Requested a conversion from torch to cupy/numpy on a non-coalesced tensor -> coalescing implicitly"
206
220
  )
207
221
  x_torch = x_torch.coalesce()
208
- data_c = xp.asarray(x_torch.values())
209
- idx_cp = xp.asarray(x_torch.indices())
222
+ data_c = _torch_to_backend(x_torch.values(), xp)
223
+ idx_cp = _torch_to_backend(x_torch.indices(), xp)
210
224
  x_cupy = xsp.coo_matrix((data_c, idx_cp), shape=x_torch.shape)
211
225
  return x_cupy
212
226
 
@@ -238,8 +252,8 @@ def c2t_coo(x_cupy: Any) -> torch.Tensor:
238
252
  >>> x.layout is torch.sparse_coo
239
253
  True
240
254
  """
241
- data_t = torch.as_tensor(x_cupy.data)
242
- row_t = torch.as_tensor(x_cupy.row)
243
- col_t = torch.as_tensor(x_cupy.col)
255
+ data_t = _backend_to_torch(x_cupy.data)
256
+ row_t = _backend_to_torch(x_cupy.row)
257
+ col_t = _backend_to_torch(x_cupy.col)
244
258
  x_torch = torch.sparse_coo_tensor(torch.stack([row_t, col_t], dim=0), data_t, x_cupy.shape)
245
259
  return x_torch