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.
- {torchsparsegradutils-0.2.2/torchsparsegradutils.egg-info → torchsparsegradutils-0.2.4}/PKG-INFO +11 -2
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/README.md +10 -1
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/docs/source/conf.py +2 -2
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/pyproject.toml +1 -1
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/__init__.py +3 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/benchmarks/benchmark_suite.py +4 -0
- torchsparsegradutils-0.2.4/torchsparsegradutils/benchmarks/sparse_bidir_logsumexp_rand.py +191 -0
- torchsparsegradutils-0.2.4/torchsparsegradutils/benchmarks/sparse_bidir_logsumexp_suitesparse.py +191 -0
- torchsparsegradutils-0.2.4/torchsparsegradutils/benchmarks/sparse_logsumexp_rand.py +175 -0
- torchsparsegradutils-0.2.4/torchsparsegradutils/benchmarks/sparse_logsumexp_suitesparse.py +176 -0
- torchsparsegradutils-0.2.4/torchsparsegradutils/sparse_logsumexp.py +496 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/sparse_solve.py +9 -2
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_doctests.py +1 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_jax_bindings.py +1 -5
- torchsparsegradutils-0.2.4/torchsparsegradutils/tests/test_sparse_bidir_logsumexp.py +299 -0
- torchsparsegradutils-0.2.4/torchsparsegradutils/tests/test_sparse_logsumexp.py +297 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_sparse_solve.py +222 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4/torchsparsegradutils.egg-info}/PKG-INFO +11 -2
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils.egg-info/SOURCES.txt +7 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/LICENSE +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/MANIFEST.in +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/setup.cfg +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/setup.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/benchmarks/__init__.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/benchmarks/batched_sparse_mm_rand.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/benchmarks/benchmark_utils.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/benchmarks/sparse_generic_solve_rand.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/benchmarks/sparse_generic_solve_suite.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/benchmarks/sparse_mm_rand.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/benchmarks/sparse_mm_suite.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/benchmarks/sparse_triangular_solve_rand.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/benchmarks/sparse_triangular_solve_suitesparse.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/benchmarks/visualize_benchmark_results.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/cupy/__init__.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/cupy/cupy_bindings.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/cupy/cupy_sparse_solve.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/distributions/__init__.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/distributions/constraints.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/distributions/sparse_multivariate_normal.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/encoders/__init__.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/encoders/pairwise_encoder.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/encoders/pairwise_voxel_encoder.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/indexed_matmul.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/jax/__init__.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/jax/jax_bindings.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/jax/jax_sparse_solve.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/sparse_lstsq.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/sparse_matmul.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/__init__.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/conftest.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_bicgstab.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_config.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_cupy_bindings.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_cupy_sparse_solve.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_dist_stats_helpers.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_distributions.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_encoders.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_indexed_matmul.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_integration_pairwise_sparse_mvn.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_jax_sparse_solve.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_linear_cg.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_lsmr.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_minres.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_params/czyx_shifts.yaml +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_params/pairwise_coo_indices.yaml +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_params/xyz_coords.yaml +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_quickstart_guide.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_random.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_sparse_lstsq.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_sparse_matmul.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_sparse_triangular_solve.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/tests/test_utils.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/utils/__init__.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/utils/bicgstab.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/utils/dist_stats_helpers.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/utils/linear_cg.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/utils/lsmr.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/utils/minres.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/utils/random_sparse.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils/utils/utils.py +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils.egg-info/dependency_links.txt +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils.egg-info/not-zip-safe +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils.egg-info/requires.txt +0 -0
- {torchsparsegradutils-0.2.2 → torchsparsegradutils-0.2.4}/torchsparsegradutils.egg-info/top_level.txt +0 -0
{torchsparsegradutils-0.2.2/torchsparsegradutils.egg-info → torchsparsegradutils-0.2.4}/PKG-INFO
RENAMED
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: torchsparsegradutils
|
|
3
|
-
Version: 0.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
|
-
[](https://pypi.org/project/torchsparsegradutils/) [](https://pypi.org/project/torchsparsegradutils/) [](https://pypi.org/project/torchsparsegradutils/)  ](https://pypi.org/project/torchsparsegradutils/) [](https://pypi.org/project/torchsparsegradutils/) [](https://pypi.org/project/torchsparsegradutils/)   [](https://github.com/cai4cai/torchsparsegradutils/actions/workflows/python-package.yml) [](https://readthedocs.org/projects/torchsparsegradutils) [](https://github.com/psf/black) [](LICENSE) [](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
|
-
[](https://pypi.org/project/torchsparsegradutils/) [](https://pypi.org/project/torchsparsegradutils/) [](https://pypi.org/project/torchsparsegradutils/)  ](https://pypi.org/project/torchsparsegradutils/) [](https://pypi.org/project/torchsparsegradutils/) [](https://pypi.org/project/torchsparsegradutils/)   [](https://github.com/cai4cai/torchsparsegradutils/actions/workflows/python-package.yml) [](https://readthedocs.org/projects/torchsparsegradutils) [](https://github.com/psf/black) [](LICENSE) [](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.
|
|
19
|
-
version = "0.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
|
|
@@ -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()
|
torchsparsegradutils-0.2.4/torchsparsegradutils/benchmarks/sparse_bidir_logsumexp_suitesparse.py
ADDED
|
@@ -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()
|