torchsparsegradutils 0.2.2__tar.gz → 0.2.4__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 (84) hide show
  1. {torchsparsegradutils-0.2.2/torchsparsegradutils.egg-info → torchsparsegradutils-0.2.4}/PKG-INFO +11 -2
  2. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/README.md +10 -1
  3. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/docs/source/conf.py +2 -2
  4. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/pyproject.toml +1 -1
  5. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/__init__.py +3 -0
  6. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/benchmarks/benchmark_suite.py +4 -0
  7. torchsparsegradutils-0.2.4/torchsparsegradutils/benchmarks/sparse_bidir_logsumexp_rand.py +191 -0
  8. torchsparsegradutils-0.2.4/torchsparsegradutils/benchmarks/sparse_bidir_logsumexp_suitesparse.py +191 -0
  9. torchsparsegradutils-0.2.4/torchsparsegradutils/benchmarks/sparse_logsumexp_rand.py +175 -0
  10. torchsparsegradutils-0.2.4/torchsparsegradutils/benchmarks/sparse_logsumexp_suitesparse.py +176 -0
  11. torchsparsegradutils-0.2.4/torchsparsegradutils/sparse_logsumexp.py +496 -0
  12. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/sparse_solve.py +9 -2
  13. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_doctests.py +1 -0
  14. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_jax_bindings.py +1 -5
  15. torchsparsegradutils-0.2.4/torchsparsegradutils/tests/test_sparse_bidir_logsumexp.py +299 -0
  16. torchsparsegradutils-0.2.4/torchsparsegradutils/tests/test_sparse_logsumexp.py +297 -0
  17. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_sparse_solve.py +222 -0
  18. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4/torchsparsegradutils.egg-info}/PKG-INFO +11 -2
  19. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils.egg-info/SOURCES.txt +7 -0
  20. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/LICENSE +0 -0
  21. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/MANIFEST.in +0 -0
  22. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/setup.cfg +0 -0
  23. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/setup.py +0 -0
  24. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/benchmarks/__init__.py +0 -0
  25. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/benchmarks/batched_sparse_mm_rand.py +0 -0
  26. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/benchmarks/benchmark_utils.py +0 -0
  27. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/benchmarks/sparse_generic_solve_rand.py +0 -0
  28. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/benchmarks/sparse_generic_solve_suite.py +0 -0
  29. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/benchmarks/sparse_mm_rand.py +0 -0
  30. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/benchmarks/sparse_mm_suite.py +0 -0
  31. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/benchmarks/sparse_triangular_solve_rand.py +0 -0
  32. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/benchmarks/sparse_triangular_solve_suitesparse.py +0 -0
  33. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/benchmarks/visualize_benchmark_results.py +0 -0
  34. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/cupy/__init__.py +0 -0
  35. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/cupy/cupy_bindings.py +0 -0
  36. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/cupy/cupy_sparse_solve.py +0 -0
  37. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/distributions/__init__.py +0 -0
  38. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/distributions/constraints.py +0 -0
  39. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/distributions/sparse_multivariate_normal.py +0 -0
  40. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/encoders/__init__.py +0 -0
  41. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/encoders/pairwise_encoder.py +0 -0
  42. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/encoders/pairwise_voxel_encoder.py +0 -0
  43. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/indexed_matmul.py +0 -0
  44. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/jax/__init__.py +0 -0
  45. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/jax/jax_bindings.py +0 -0
  46. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/jax/jax_sparse_solve.py +0 -0
  47. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/sparse_lstsq.py +0 -0
  48. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/sparse_matmul.py +0 -0
  49. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/__init__.py +0 -0
  50. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/conftest.py +0 -0
  51. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_bicgstab.py +0 -0
  52. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_config.py +0 -0
  53. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_cupy_bindings.py +0 -0
  54. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_cupy_sparse_solve.py +0 -0
  55. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_dist_stats_helpers.py +0 -0
  56. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_distributions.py +0 -0
  57. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_encoders.py +0 -0
  58. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_indexed_matmul.py +0 -0
  59. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_integration_pairwise_sparse_mvn.py +0 -0
  60. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_jax_sparse_solve.py +0 -0
  61. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_linear_cg.py +0 -0
  62. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_lsmr.py +0 -0
  63. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_minres.py +0 -0
  64. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_params/czyx_shifts.yaml +0 -0
  65. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_params/pairwise_coo_indices.yaml +0 -0
  66. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_params/xyz_coords.yaml +0 -0
  67. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_quickstart_guide.py +0 -0
  68. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_random.py +0 -0
  69. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_sparse_lstsq.py +0 -0
  70. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_sparse_matmul.py +0 -0
  71. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_sparse_triangular_solve.py +0 -0
  72. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_utils.py +0 -0
  73. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/utils/__init__.py +0 -0
  74. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/utils/bicgstab.py +0 -0
  75. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/utils/dist_stats_helpers.py +0 -0
  76. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/utils/linear_cg.py +0 -0
  77. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/utils/lsmr.py +0 -0
  78. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/utils/minres.py +0 -0
  79. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/utils/random_sparse.py +0 -0
  80. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/utils/utils.py +0 -0
  81. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils.egg-info/dependency_links.txt +0 -0
  82. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils.egg-info/not-zip-safe +0 -0
  83. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils.egg-info/requires.txt +0 -0
  84. {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils.egg-info/top_level.txt +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: torchsparsegradutils
3
- Version: 0.2.2
3
+ Version: 0.2.4
4
4
  Summary: A collection of utility functions to work with PyTorch sparse tensors
5
5
  Author-email: CAI4CAI research group <contact@cai4cai.uk>
6
6
  License-Expression: Apache-2.0
@@ -47,7 +47,7 @@ Dynamic: license-file
47
47
 
48
48
  # torchsparsegradutils: Sparsity-preserving gradient utility tools for PyTorch
49
49
 
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)
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.13 / nightly](https://img.shields.io/badge/Tested-2.5%20|%202.13%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)
51
51
 
52
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.
53
53
 
@@ -61,6 +61,15 @@ A comprehensive collection of utility functions to work with PyTorch sparse tens
61
61
  - Workaround for [PyTorch issue #41128](https://github.com/pytorch/pytorch/issues/41128)
62
62
  - Supports both COO and CSR formats with optional batching
63
63
 
64
+ **Numerically-Stable Sparse Reductions**
65
+ - `sparse_logsumexp`: Sparse-aware `log-sum-exp` reduction mirroring `torch.logsumexp`
66
+ - Operates directly on the nonzero values (no dense materialisation), with a numerically stable max-shift
67
+ - Supports COO/CSR/CSC layouts (unbatched 2-D and batched 3-D) and an `include_zeros` flag for structural-zero semantics
68
+ - Fills the gap of [PyTorch issue #31394](https://github.com/pytorch/pytorch/issues/31394) (no native `scatter_logsumexp`)
69
+ - `sparse_bidir_logsumexp`: Row- and column-wise `log-sum-exp` simultaneously in a single traversal
70
+ - Fuses the two `sparse_logsumexp(dim=0)` / `sparse_logsumexp(dim=1)` passes into one batched scatter, sharing the index extraction and autograd graph
71
+ - Returns `tuple`, `padded`, or `nested` output layouts
72
+
64
73
  **Sparse Linear System Solvers**
65
74
  - `sparse_triangular_solve`: Sparse triangular solver with batch support
66
75
  - Discussion reference: [PyTorch issue #87358](https://github.com/pytorch/pytorch/issues/87358)
@@ -1,6 +1,6 @@
1
1
  # torchsparsegradutils: Sparsity-preserving gradient utility tools for PyTorch
2
2
 
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)
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.13 / nightly](https://img.shields.io/badge/Tested-2.5%20|%202.13%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
 
@@ -14,6 +14,15 @@ A comprehensive collection of utility functions to work with PyTorch sparse tens
14
14
  - Workaround for [PyTorch issue #41128](https://github.com/pytorch/pytorch/issues/41128)
15
15
  - Supports both COO and CSR formats with optional batching
16
16
 
17
+ **Numerically-Stable Sparse Reductions**
18
+ - `sparse_logsumexp`: Sparse-aware `log-sum-exp` reduction mirroring `torch.logsumexp`
19
+ - Operates directly on the nonzero values (no dense materialisation), with a numerically stable max-shift
20
+ - Supports COO/CSR/CSC layouts (unbatched 2-D and batched 3-D) and an `include_zeros` flag for structural-zero semantics
21
+ - Fills the gap of [PyTorch issue #31394](https://github.com/pytorch/pytorch/issues/31394) (no native `scatter_logsumexp`)
22
+ - `sparse_bidir_logsumexp`: Row- and column-wise `log-sum-exp` simultaneously in a single traversal
23
+ - Fuses the two `sparse_logsumexp(dim=0)` / `sparse_logsumexp(dim=1)` passes into one batched scatter, sharing the index extraction and autograd graph
24
+ - Returns `tuple`, `padded`, or `nested` output layouts
25
+
17
26
  **Sparse Linear System Solvers**
18
27
  - `sparse_triangular_solve`: Sparse triangular solver with batch support
19
28
  - Discussion reference: [PyTorch issue #87358](https://github.com/pytorch/pytorch/issues/87358)
@@ -15,8 +15,8 @@ sys.path.insert(0, os.path.abspath("../../"))
15
15
  project = "torchsparsegradutils"
16
16
  copyright = "2026, CAI4CAI research group"
17
17
  author = "CAI4CAI research group"
18
- release = "0.2.2"
19
- version = "0.2.2"
18
+ release = "0.2.4"
19
+ version = "0.2.4"
20
20
 
21
21
  # -- General configuration ---------------------------------------------------
22
22
  # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "torchsparsegradutils"
7
- version = "0.2.2"
7
+ version = "0.2.4"
8
8
  description = "A collection of utility functions to work with PyTorch sparse tensors"
9
9
  readme = "README.md"
10
10
  requires-python = ">=3.10"
@@ -1,4 +1,5 @@
1
1
  from .indexed_matmul import gather_mm, segment_mm
2
+ from .sparse_logsumexp import sparse_bidir_logsumexp, sparse_logsumexp
2
3
  from .sparse_lstsq import sparse_generic_lstsq
3
4
  from .sparse_matmul import sparse_mm
4
5
  from .sparse_solve import sparse_generic_solve, sparse_triangular_solve
@@ -10,4 +11,6 @@ __all__ = [
10
11
  "sparse_triangular_solve",
11
12
  "sparse_generic_solve",
12
13
  "sparse_generic_lstsq",
14
+ "sparse_logsumexp",
15
+ "sparse_bidir_logsumexp",
13
16
  ]
@@ -19,6 +19,10 @@ BENCHMARK_SCRIPTS = [
19
19
  "sparse_triangular_solve_rand.py",
20
20
  "sparse_triangular_solve_suitesparse.py",
21
21
  "batched_sparse_mm_rand.py",
22
+ "sparse_logsumexp_rand.py",
23
+ "sparse_logsumexp_suitesparse.py",
24
+ "sparse_bidir_logsumexp_rand.py",
25
+ "sparse_bidir_logsumexp_suitesparse.py",
22
26
  ]
23
27
 
24
28
 
@@ -0,0 +1,191 @@
1
+ #!/usr/bin/env python3
2
+ """
3
+ Sparse Bidirectional Log-Sum-Exp Benchmark - Random Matrices
4
+
5
+ Benchmarks ``sparse_bidir_logsumexp`` (both row- and column-wise reductions in a
6
+ single traversal) against the two-call baseline
7
+ ``(sparse_logsumexp(A, dim=0), sparse_logsumexp(A, dim=1))`` on randomly generated
8
+ sparse matrices of various sizes, sparsity patterns, layouts and dtypes. Forward
9
+ and backward time and peak memory are measured; the baseline is what the single-
10
+ pass primitive exists to beat. Both algorithms are driven through a single tensor
11
+ (``output_layout="padded"`` / ``torch.cat``) so ``measure_op``'s ``.sum().backward()``
12
+ exercises both reductions.
13
+ """
14
+
15
+ import os
16
+ import sys
17
+
18
+ # Add the parent directory to sys.path to allow importing torchsparsegradutils
19
+ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
20
+
21
+ import numpy as np
22
+ import torch
23
+ from benchmark_utils import (
24
+ measure_op,
25
+ print_benchmark_header,
26
+ print_result_row,
27
+ print_results_table_header,
28
+ save_benchmark_results,
29
+ )
30
+ from tqdm import tqdm
31
+
32
+ from torchsparsegradutils import sparse_bidir_logsumexp, sparse_logsumexp
33
+ from torchsparsegradutils.utils import rand_sparse
34
+
35
+ REPEATS = 100
36
+ WARMUP_RUNS = 10
37
+
38
+ # Only run on CUDA
39
+ device = torch.device("cuda")
40
+ assert torch.cuda.is_available(), "This benchmark requires a CUDA GPU"
41
+
42
+ # problem sizes: (label, N, M, nnz)
43
+ SIZES = [
44
+ ("small", 2**10, 2**10, 2**12),
45
+ ("medium", 2**12, 2**12, 2**14),
46
+ ("large", 2**14, 2**14, 2**16),
47
+ ("xlarge", 2**16, 2**16, 2**18),
48
+ ("million", 2**20, 2**20, 2**22),
49
+ ]
50
+
51
+ INDEX_DTYPES = [torch.int32, torch.int64]
52
+ VALUE_DTYPES = [torch.float32, torch.float64]
53
+ LAYOUTS = [torch.sparse_coo, torch.sparse_csr]
54
+
55
+ # Both ops return a single tensor so measure_op's out.sum().backward() drives both
56
+ # reductions. B is an unused placeholder for measure_op's binary signature.
57
+ ALGORITHMS = [
58
+ ("sparse_bidir_logsumexp", lambda A, B: sparse_bidir_logsumexp(A, output_layout="padded")),
59
+ ("two_call_baseline", lambda A, B: torch.cat([sparse_logsumexp(A, dim=0), sparse_logsumexp(A, dim=1)])),
60
+ ]
61
+
62
+
63
+ def run_sparse_bidir_logsumexp_benchmark():
64
+ """Run the sparse bidirectional log-sum-exp benchmark suite."""
65
+
66
+ print_benchmark_header("Sparse Bidirectional Log-Sum-Exp Benchmark - Random Matrices")
67
+
68
+ records = []
69
+
70
+ for size_label, N, M, nnz in tqdm(SIZES, desc="Problem sizes"):
71
+ print(f"\n🔍 Testing size: {size_label} (N={N}, M={M}, nnz={nnz})")
72
+
73
+ A_shape = (N, M)
74
+ B = torch.zeros(1, device=device) # unused placeholder for measure_op's binary signature
75
+
76
+ for idx_dt in tqdm(INDEX_DTYPES, desc="Index dtypes", leave=False):
77
+ for val_dt in tqdm(VALUE_DTYPES, desc="Value dtypes", leave=False):
78
+ for layout in tqdm(LAYOUTS, desc="Layouts", leave=False):
79
+ layout_name = "coo" if layout == torch.sparse_coo else "csr"
80
+ print(f"\n 📊 Configuration: idx_dtype={idx_dt}, val_dtype={val_dt}, layout={layout_name}")
81
+
82
+ # COO indices are always int64 -> skip int32 rather than mislabel a duplicate.
83
+ if layout == torch.sparse_coo and idx_dt != torch.int64:
84
+ print(f" ⏭ skipping {layout_name} + {idx_dt} (COO indices are always int64)")
85
+ continue
86
+
87
+ # Generate once per config, outside the try so generation errors surface here.
88
+ A_sparse = rand_sparse(
89
+ A_shape, nnz, layout, indices_dtype=idx_dt, values_dtype=val_dt, device=device
90
+ )
91
+ actual_idx_dt = (
92
+ A_sparse.col_indices().dtype
93
+ if layout == torch.sparse_csr
94
+ else A_sparse.coalesce().indices().dtype
95
+ )
96
+ assert actual_idx_dt == idx_dt, f"index dtype mislabel: requested {idx_dt}, built {actual_idx_dt}"
97
+
98
+ print_results_table_header()
99
+
100
+ for alg_name, alg_fn in ALGORITHMS:
101
+ try:
102
+ print(f" 🧮 Testing {alg_name} ({layout_name})...")
103
+
104
+ (
105
+ t_fwd,
106
+ std_fwd,
107
+ mem_fwd,
108
+ std_mem_fwd,
109
+ t_bwd,
110
+ std_bwd,
111
+ mem_bwd,
112
+ std_mem_bwd,
113
+ ) = measure_op(
114
+ alg_fn,
115
+ A_sparse,
116
+ B,
117
+ repeats=REPEATS,
118
+ device=device,
119
+ desc=f"{alg_name} ({layout_name})",
120
+ warmup_runs=WARMUP_RUNS,
121
+ remove_outliers=True,
122
+ )
123
+
124
+ print_result_row(
125
+ f"{alg_name} ({layout_name})",
126
+ (N, M),
127
+ t_fwd,
128
+ std_fwd,
129
+ mem_fwd,
130
+ std_mem_fwd,
131
+ t_bwd,
132
+ std_bwd,
133
+ mem_bwd,
134
+ std_mem_bwd,
135
+ )
136
+
137
+ records.append(
138
+ {
139
+ "size": size_label,
140
+ "layout": layout_name,
141
+ "algo": alg_name,
142
+ "index_dt": str(idx_dt).split(".")[-1],
143
+ "value_dt": str(val_dt).split(".")[-1],
144
+ "N": N,
145
+ "M": M,
146
+ "nnz": nnz,
147
+ "fwd_time_us": t_fwd,
148
+ "fwd_time_std_us": std_fwd,
149
+ "fwd_mem_MB": mem_fwd,
150
+ "fwd_mem_std_MB": std_mem_fwd,
151
+ "bwd_time_us": t_bwd,
152
+ "bwd_time_std_us": std_bwd,
153
+ "bwd_mem_MB": mem_bwd,
154
+ "bwd_mem_std_MB": std_mem_bwd,
155
+ }
156
+ )
157
+
158
+ except Exception as e:
159
+ print(f" ❌ {alg_name} ({layout_name}) failed: {e}")
160
+
161
+ records.append(
162
+ {
163
+ "size": size_label,
164
+ "layout": layout_name,
165
+ "algo": alg_name,
166
+ "index_dt": str(idx_dt).split(".")[-1],
167
+ "value_dt": str(val_dt).split(".")[-1],
168
+ "N": N,
169
+ "M": M,
170
+ "nnz": nnz,
171
+ "fwd_time_us": np.nan,
172
+ "fwd_time_std_us": np.nan,
173
+ "fwd_mem_MB": np.nan,
174
+ "fwd_mem_std_MB": np.nan,
175
+ "bwd_time_us": np.nan,
176
+ "bwd_time_std_us": np.nan,
177
+ "bwd_mem_MB": np.nan,
178
+ "bwd_mem_std_MB": np.nan,
179
+ "error": str(e),
180
+ }
181
+ )
182
+
183
+ # Save results
184
+ if records:
185
+ save_benchmark_results(records, "sparse_bidir_logsumexp_rand")
186
+
187
+ print("\n✅ Sparse bidirectional log-sum-exp benchmark completed!")
188
+
189
+
190
+ if __name__ == "__main__":
191
+ run_sparse_bidir_logsumexp_benchmark()
@@ -0,0 +1,191 @@
1
+ #!/usr/bin/env python3
2
+ """
3
+ Sparse Bidirectional Log-Sum-Exp Benchmark - SuiteSparse Collection
4
+
5
+ Benchmarks ``sparse_bidir_logsumexp`` (both row- and column-wise reductions in a
6
+ single traversal) against the two-call baseline
7
+ ``(sparse_logsumexp(A, dim=0), sparse_logsumexp(A, dim=1))`` on real matrices from
8
+ the SuiteSparse Matrix Collection (Rothberg/cfd2 ~123k x 123k, Williams/webbase-1M
9
+ ~1M x 1M). Forward and backward time and peak memory are measured. Both algorithms
10
+ are driven through a single tensor (``output_layout="padded"`` / ``torch.cat``) so
11
+ ``measure_op``'s ``.sum().backward()`` exercises both reductions.
12
+ """
13
+
14
+ import os
15
+ import sys
16
+
17
+ # Add the parent directory to sys.path to allow importing torchsparsegradutils
18
+ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
19
+
20
+ import numpy as np
21
+ import torch
22
+ from benchmark_utils import (
23
+ load_mat_from_suitesparse_collection,
24
+ measure_op,
25
+ print_benchmark_header,
26
+ print_result_row,
27
+ print_results_table_header,
28
+ save_benchmark_results,
29
+ )
30
+ from tqdm import tqdm
31
+
32
+ from torchsparsegradutils import sparse_bidir_logsumexp, sparse_logsumexp
33
+
34
+ REPEATS = 100
35
+ WARMUP_RUNS = 10
36
+
37
+ # Only run on CUDA
38
+ device = torch.device("cuda")
39
+ assert torch.cuda.is_available(), "This benchmark requires a CUDA GPU"
40
+
41
+ INDEX_DTYPES = [torch.int32, torch.int64]
42
+ VALUE_DTYPES = [torch.float32, torch.float64]
43
+ LAYOUTS = [torch.sparse_coo, torch.sparse_csr]
44
+
45
+ # Both ops return a single tensor so measure_op's out.sum().backward() drives both
46
+ # reductions. B is an unused placeholder for measure_op's binary signature.
47
+ ALGORITHMS = [
48
+ ("sparse_bidir_logsumexp", lambda A, B: sparse_bidir_logsumexp(A, output_layout="padded")),
49
+ ("two_call_baseline", lambda A, B: torch.cat([sparse_logsumexp(A, dim=0), sparse_logsumexp(A, dim=1)])),
50
+ ]
51
+
52
+ MATRICES = [
53
+ ("Rothberg", "cfd2"), # 123k x 123k
54
+ ("Williams", "webbase-1M"), # 1M x 1M
55
+ ]
56
+
57
+
58
+ def run_sparse_bidir_logsumexp_suitesparse_benchmark():
59
+ """Run the sparse bidirectional log-sum-exp benchmark suite with SuiteSparse matrices."""
60
+
61
+ print_benchmark_header("Sparse Bidirectional Log-Sum-Exp Benchmark - SuiteSparse Collection")
62
+
63
+ records = []
64
+ B = torch.zeros(1, device=device) # unused placeholder for measure_op's binary signature
65
+
66
+ for dirname, matname in tqdm(MATRICES, desc="Matrices"):
67
+ A_np_coo = load_mat_from_suitesparse_collection(dirname, matname)
68
+ matrix_name = f"{dirname}/{matname}"
69
+ N, M = A_np_coo.shape
70
+ nnz = A_np_coo.nnz
71
+ print(f"📊 Matrix: {matrix_name}, shape={A_np_coo.shape}, nnz={nnz}")
72
+
73
+ for idx_dt in tqdm(INDEX_DTYPES, desc="Index dtypes", leave=False):
74
+ for val_dt in tqdm(VALUE_DTYPES, desc="Value dtypes", leave=False):
75
+ for layout in tqdm(LAYOUTS, desc="Layouts", leave=False):
76
+ layout_name = "coo" if layout == torch.sparse_coo else "csr"
77
+ print(f"\n🔍 Testing configuration: idx_dtype={idx_dt}, val_dtype={val_dt}, layout={layout_name}")
78
+
79
+ # COO indices are always int64 -> skip int32 rather than mislabel a duplicate.
80
+ if layout == torch.sparse_coo and idx_dt != torch.int64:
81
+ print(f" ⏭ skipping {layout_name} + {idx_dt} (COO indices are always int64)")
82
+ continue
83
+
84
+ values = torch.tensor(A_np_coo.data, dtype=val_dt, device=device)
85
+ indices = torch.tensor([A_np_coo.row, A_np_coo.col], dtype=torch.int64, device=device)
86
+ A_coo = torch.sparse_coo_tensor(
87
+ indices, values, A_np_coo.shape, dtype=val_dt, device=device
88
+ ).coalesce()
89
+ if layout == torch.sparse_csr:
90
+ # to_sparse_csr() gives int64 -> rebuild with the requested dtype so the label is honest.
91
+ csr = A_coo.to_sparse_csr()
92
+ A_sparse = torch.sparse_csr_tensor(
93
+ csr.crow_indices().to(idx_dt), csr.col_indices().to(idx_dt), csr.values(), csr.shape
94
+ )
95
+ assert A_sparse.col_indices().dtype == idx_dt
96
+ else:
97
+ A_sparse = A_coo
98
+
99
+ print_results_table_header()
100
+
101
+ for alg_name, alg_fn in ALGORITHMS:
102
+ try:
103
+ print(f" 🧮 Testing {alg_name} ({layout_name})...")
104
+
105
+ (
106
+ t_fwd,
107
+ std_fwd,
108
+ mem_fwd,
109
+ std_mem_fwd,
110
+ t_bwd,
111
+ std_bwd,
112
+ mem_bwd,
113
+ std_mem_bwd,
114
+ ) = measure_op(
115
+ alg_fn,
116
+ A_sparse,
117
+ B,
118
+ repeats=REPEATS,
119
+ device=device,
120
+ desc=f"{alg_name} ({layout_name})",
121
+ warmup_runs=WARMUP_RUNS,
122
+ remove_outliers=True,
123
+ )
124
+
125
+ print_result_row(
126
+ f"{alg_name} ({layout_name})",
127
+ (N, M),
128
+ t_fwd,
129
+ std_fwd,
130
+ mem_fwd,
131
+ std_mem_fwd,
132
+ t_bwd,
133
+ std_bwd,
134
+ mem_bwd,
135
+ std_mem_bwd,
136
+ )
137
+
138
+ records.append(
139
+ {
140
+ "matrix": matrix_name,
141
+ "N": N,
142
+ "M": M,
143
+ "nnz": nnz,
144
+ "index_dt": str(idx_dt).split(".")[-1],
145
+ "value_dt": str(val_dt).split(".")[-1],
146
+ "layout": layout_name,
147
+ "algo": alg_name,
148
+ "fwd_time_us": t_fwd,
149
+ "fwd_time_std_us": std_fwd,
150
+ "fwd_mem_MB": mem_fwd,
151
+ "fwd_mem_std_MB": std_mem_fwd,
152
+ "bwd_time_us": t_bwd,
153
+ "bwd_time_std_us": std_bwd,
154
+ "bwd_mem_MB": mem_bwd,
155
+ "bwd_mem_std_MB": std_mem_bwd,
156
+ }
157
+ )
158
+
159
+ except Exception as e:
160
+ print(f" ❌ {alg_name} ({layout_name}) failed: {e}")
161
+
162
+ records.append(
163
+ {
164
+ "matrix": matrix_name,
165
+ "N": N,
166
+ "M": M,
167
+ "nnz": nnz,
168
+ "index_dt": str(idx_dt).split(".")[-1],
169
+ "value_dt": str(val_dt).split(".")[-1],
170
+ "layout": layout_name,
171
+ "algo": alg_name,
172
+ "fwd_time_us": np.nan,
173
+ "fwd_time_std_us": np.nan,
174
+ "fwd_mem_MB": np.nan,
175
+ "fwd_mem_std_MB": np.nan,
176
+ "bwd_time_us": np.nan,
177
+ "bwd_time_std_us": np.nan,
178
+ "bwd_mem_MB": np.nan,
179
+ "bwd_mem_std_MB": np.nan,
180
+ "error": str(e),
181
+ }
182
+ )
183
+
184
+ if records:
185
+ save_benchmark_results(records, "sparse_bidir_logsumexp_suitesparse")
186
+
187
+ print("\n✅ Sparse bidirectional log-sum-exp SuiteSparse benchmark completed!")
188
+
189
+
190
+ if __name__ == "__main__":
191
+ run_sparse_bidir_logsumexp_suitesparse_benchmark()