Add sparse_logsumexp: numerically stable sparse log-sum-exp reduction (#87)

* feat: add sparse_logsumexp reduction

Numerically-stable log-sum-exp over the nonzero values of a 2-D sparse
tensor (COO/CSR/CSC), mirroring torch.logsumexp. Reads CSR/CSC indices
directly (no COO coalesce) and reuses the compressed-axis nnz counts to
avoid a redundant bincount. The max-shift is detached so gradients equal
softmax over the participating entries. An include_zeros flag selects
whether structural zeros contribute exp(0)=1 or are treated as -inf.

* test: cover sparse_logsumexp against dense reference

Parametrized over COO/CSR layouts, devices, dtypes, all dim modes and
both include_zeros settings; cross-checks against torch.logsumexp on the
densified tensor. Adds keepdim, CSC, empty row/column, error-path and
gradient tests, and registers the module for doctest execution.

* bench: add sparse_logsumexp benchmark

Compares sparse_logsumexp against a dense torch.logsumexp baseline and a
segment_reduce variant across sizes/densities, on CPU and CUDA, with
graceful handling of dense out-of-memory at large sizes.

* docs: document sparse_logsumexp in API reference and README

* feat: support batched 3-D sparse_logsumexp

Extend the reduction to batched 3-D (batch, rows, cols) inputs, matching
the batched convention of sparse_mm and the solvers. Each slice is reduced
independently by folding the batch index into the scatter segment id, so a
single scatter handles the whole batch; the batch axis itself is not
reducible. Batched inputs go through COO (batched CSR/CSC require equal nnz
per slice in PyTorch). Adds batched tests across dims/include_zeros/keepdim.

* refactor: rename segment->scatter in logsumexp helper

The helper reduces via Tensor.scatter_reduce_ with an arbitrary per-value
output index, so 'segment' (which implies contiguous runs, as in
torch.segment_reduce) was misleading. Rename to _scatter_logsumexp /
scatter_index and note in the docstring that group members need not be
contiguous.

* refactor: use torch.logsumexp directly for scalar reduction

For dim=[0,1] every value falls in one group, so torch.logsumexp does the
reduction directly. The structural-zero mass (n_zeros entries of exp(0)=1)
equals a single entry of value log(n_zeros), appended so the reduction stays
one numerically stable call whose gradient still accounts for the zeros.

* docs: note why scatter is used over a segmented reduction for CSR/CSC

* docs: clarify 'slice' means one batch matrix input[k]

* bench: cover >1M-dim and SuiteSparse matrices for sparse_logsumexp

Extend the random benchmark to 2**16 and 2**20 (>1M rows/cols) where the dense
baseline OOMs, and add sparse_logsumexp_suitesparse.py running on Rothberg/cfd2
(123k x 123k) and Williams/webbase-1M (1M x 1M). Register both in
benchmark_suite.py and track their result CSVs.

* fix: correct stability shift for structural zeros in sparse_logsumexp

With include_zeros=True the structural zeros (value 0) participate in the
reduction, but the max-shift was computed only from the explicit values. When
every explicit value in a group is negative this overflowed (exp(-shift) -> inf),
and a fully dense group (no structural zeros) hit 0 * inf -> nan. Include 0 in
the per-group max where structural zeros are present, and contribute the zero
term only where the count is nonzero. Adds regression tests.

Reported-by: theo-barfoot

* fix: handle +inf values in sparse_logsumexp shift

An explicit +inf made the per-group max +inf, so value - shift computed
inf - inf = nan. Replace any non-finite shift (both +inf and the empty-group
-inf) with 0, so a +inf value correctly yields +inf. Adds regression test.

Reported-by: theo-barfoot

* fix: validate dim before normalising in sparse_logsumexp

dim was reduced modulo ndim before validation, so out-of-range dims silently
wrapped (2 -> 0), a set() dropped duplicates ([0, 0] -> [0]), and an empty
sequence fell through to a wrong path. Validate raw dims first: IndexError for
out-of-range, RuntimeError for empty or repeated dims, matching torch.logsumexp.
Adds regression tests; valid negative dims are unaffected.

Reported-by: theo-barfoot

* test: parametrize sparse_logsumexp tests over index dtype

Add an index_dtype fixture (int32/int64) matching the repo convention and thread
it through the dense-reference test so the CSR/CSC index paths are covered for
both dtypes. COO indices are int64-only in PyTorch, so the fixture varies only
the compressed layouts.

Suggested-by: theo-barfoot

* refactor: guard coalesce with is_coalesced in sparse_logsumexp

Skip the redundant coalesce() call when the COO input is already
coalesced, matching the guard pattern used elsewhere in the repo and at
the CSR/CSC fallback in this same file (line 138).

Suggested-by: tvercaut

Author: aymuos15

README.md

@@ -14,6 +14,12 @@ A comprehensive collection of utility functions to work with PyTorch sparse tens
 - Workaround for [PyTorch issue #41128](https://github.com/pytorch/pytorch/issues/41128)
 - Supports both COO and CSR formats with optional batching
 
+**Numerically-Stable Sparse Reductions**
+- `sparse_logsumexp`: Sparse-aware `log-sum-exp` reduction mirroring `torch.logsumexp`
+- Operates directly on the nonzero values (no dense materialisation), with a numerically stable max-shift
+- Supports COO/CSR/CSC layouts (unbatched 2-D and batched 3-D) and an `include_zeros` flag for structural-zero semantics
+- Fills the gap of [PyTorch issue #31394](https://github.com/pytorch/pytorch/issues/31394) (no native `scatter_logsumexp`)
+
 **Sparse Linear System Solvers**
 - `sparse_triangular_solve`: Sparse triangular solver with batch support
   -  Discussion reference: [PyTorch issue #87358](https://github.com/pytorch/pytorch/issues/87358)

docs/source/api/core.rst

@@ -25,6 +25,13 @@ Indexed Matrix Multiplication
 
 .. autofunction:: segment_mm
 
+Sparse Reductions
+-----------------
+
+.. currentmodule:: torchsparsegradutils.sparse_logsumexp
+
+.. autofunction:: sparse_logsumexp
+
 Sparse Linear Solvers
 ----------------------
 

torchsparsegradutils/__init__.py

@@ -1,4 +1,5 @@
 from .indexed_matmul import gather_mm, segment_mm
+from .sparse_logsumexp import sparse_logsumexp
 from .sparse_lstsq import sparse_generic_lstsq
 from .sparse_matmul import sparse_mm
 from .sparse_solve import sparse_generic_solve, sparse_triangular_solve
@@ -10,4 +11,5 @@
     "sparse_triangular_solve",
     "sparse_generic_solve",
     "sparse_generic_lstsq",
+    "sparse_logsumexp",
 ]

torchsparsegradutils/benchmarks/benchmark_suite.py

@@ -19,6 +19,8 @@
     "sparse_triangular_solve_rand.py",
     "sparse_triangular_solve_suitesparse.py",
     "batched_sparse_mm_rand.py",
+    "sparse_logsumexp_rand.py",
+    "sparse_logsumexp_suitesparse.py",
 ]
 
 

torchsparsegradutils/benchmarks/results/sparse_logsumexp_rand_results.csv

@@ -0,0 +1,33 @@
+matrix,N,M,nnz,index_dt,value_dt,layout,algo,fwd_time_us,fwd_time_std_us,fwd_mem_MB,fwd_mem_std_MB,bwd_time_us,bwd_time_std_us,bwd_mem_MB,bwd_mem_std_MB,python_version,pytorch_version,cuda_version,cuda_device_name
+Rothberg/cfd2,123440,123440,3087898,int32,float32,coo,sparse_logsumexp,418.5277514985388,4.05087693996108,192.20275200000003,2.842170943040401e-14,734.4205446912898,114.7157104775913,201.82272000000003,2.842170943040401e-14,3.12.3,2.11.0+cu128,12.8,NVIDIA GeForce RTX 5090
+Rothberg/cfd2,123440,123440,3087898,int32,float32,coo,dense.logsumexp,,,58132.48,,,,,,3.12.3,2.11.0+cu128,12.8,NVIDIA GeForce RTX 5090
+Rothberg/cfd2,123440,123440,3087898,int32,float32,csr,sparse_logsumexp,501.66498763149883,7.649678428248083,194.17804800000005,5.684341886080802e-14,770.3902083449066,107.41563711558506,194.17804800000005,5.684341886080802e-14,3.12.3,2.11.0+cu128,12.8,NVIDIA GeForce RTX 5090
+Rothberg/cfd2,123440,123440,3087898,int32,float32,csr,dense.logsumexp,,,58132.48,,,,,,3.12.3,2.11.0+cu128,12.8,NVIDIA GeForce RTX 5090
+Rothberg/cfd2,123440,123440,3087898,int32,float64,coo,sparse_logsumexp,507.07297932897364,4.405871818295325,257.58515200000005,5.684341886080802e-14,825.7217025170482,105.87526976135317,277.81376000000006,5.684341886080802e-14,3.12.3,2.11.0+cu128,12.8,NVIDIA GeForce RTX 5090
+Rothberg/cfd2,123440,123440,3087898,int32,float64,coo,dense.logsumexp,,,116254.72,,,,,,3.12.3,2.11.0+cu128,12.8,NVIDIA GeForce RTX 5090
+Rothberg/cfd2,123440,123440,3087898,int32,float64,csr,sparse_logsumexp,564.1494516136223,6.06369109543913,259.56044800000006,5.684341886080802e-14,859.5567787415348,88.63302570224444,259.56044800000006,5.684341886080802e-14,3.12.3,2.11.0+cu128,12.8,NVIDIA GeForce RTX 5090
+Rothberg/cfd2,123440,123440,3087898,int32,float64,csr,dense.logsumexp,,,116254.72,,,,,,3.12.3,2.11.0+cu128,12.8,NVIDIA GeForce RTX 5090
+Rothberg/cfd2,123440,123440,3087898,int64,float32,coo,sparse_logsumexp,416.6980009661832,9.744676566465959,217.59999999999988,1.1368683772161603e-13,711.014287009341,116.1010905243,227.21996799999997,2.842170943040401e-14,3.12.3,2.11.0+cu128,12.8,NVIDIA GeForce RTX 5090
+Rothberg/cfd2,123440,123440,3087898,int64,float32,coo,dense.logsumexp,,,58132.48,,,,,,3.12.3,2.11.0+cu128,12.8,NVIDIA GeForce RTX 5090
+Rothberg/cfd2,123440,123440,3087898,int64,float32,csr,sparse_logsumexp,451.92705289719953,6.678284499637928,218.651136,0.0,710.8004911060561,122.89434886626724,218.651136,0.0,3.12.3,2.11.0+cu128,12.8,NVIDIA GeForce RTX 5090
+Rothberg/cfd2,123440,123440,3087898,int64,float32,csr,dense.logsumexp,,,58132.48,,,,,,3.12.3,2.11.0+cu128,12.8,NVIDIA GeForce RTX 5090
+Rothberg/cfd2,123440,123440,3087898,int64,float64,coo,sparse_logsumexp,508.09487109528044,5.7692700070140654,282.2886399999999,1.1368683772161603e-13,844.6509094974684,101.03218358009538,302.51724800000005,5.684341886080802e-14,3.12.3,2.11.0+cu128,12.8,NVIDIA GeForce RTX 5090
+Rothberg/cfd2,123440,123440,3087898,int64,float64,coo,dense.logsumexp,,,116254.72,,,,,,3.12.3,2.11.0+cu128,12.8,NVIDIA GeForce RTX 5090
+Rothberg/cfd2,123440,123440,3087898,int64,float64,csr,sparse_logsumexp,575.4054993461898,7.326336114688565,284.72627199999994,5.684341886080802e-14,874.650359187352,85.90047695500046,284.72627199999994,5.684341886080802e-14,3.12.3,2.11.0+cu128,12.8,NVIDIA GeForce RTX 5090
+Rothberg/cfd2,123440,123440,3087898,int64,float64,csr,dense.logsumexp,,,116254.72,,,,,,3.12.3,2.11.0+cu128,12.8,NVIDIA GeForce RTX 5090
+Williams/webbase-1M,1000005,1000005,3105536,int32,float32,coo,sparse_logsumexp,481.70156716220475,4.23519804752169,220.52249600000007,8.526512829121202e-14,557.1488393198183,67.37703219108285,220.52249600000007,8.526512829121202e-14,3.12.3,2.11.0+cu128,12.8,NVIDIA GeForce RTX 5090
+Williams/webbase-1M,1000005,1000005,3105536,int32,float32,coo,dense.logsumexp,,,3814737.92,,,,,,3.12.3,2.11.0+cu128,12.8,NVIDIA GeForce RTX 5090
+Williams/webbase-1M,1000005,1000005,3105536,int32,float32,csr,sparse_logsumexp,552.8516274817447,3.7897933778308235,239.30112000000005,5.684341886080802e-14,676.7854728662458,65.88588450101828,239.30112000000005,5.684341886080802e-14,3.12.3,2.11.0+cu128,12.8,NVIDIA GeForce RTX 5090
+Williams/webbase-1M,1000005,1000005,3105536,int32,float32,csr,dense.logsumexp,,,3814737.92,,,,,,3.12.3,2.11.0+cu128,12.8,NVIDIA GeForce RTX 5090
+Williams/webbase-1M,1000005,1000005,3105536,int32,float64,coo,sparse_logsumexp,626.365046395841,5.4730458120651555,307.17798400000004,5.684341886080802e-14,881.5396050899054,84.99292983227919,307.17798400000004,5.684341886080802e-14,3.12.3,2.11.0+cu128,12.8,NVIDIA GeForce RTX 5090
+Williams/webbase-1M,1000005,1000005,3105536,int32,float64,coo,dense.logsumexp,,,7629475.84,,,,,,3.12.3,2.11.0+cu128,12.8,NVIDIA GeForce RTX 5090
+Williams/webbase-1M,1000005,1000005,3105536,int32,float64,csr,sparse_logsumexp,669.604426245011,5.613909730105387,324.66483199999993,5.684341886080802e-14,846.7558345728321,49.59094487783846,324.66483199999993,5.684341886080802e-14,3.12.3,2.11.0+cu128,12.8,NVIDIA GeForce RTX 5090
+Williams/webbase-1M,1000005,1000005,3105536,int32,float64,csr,dense.logsumexp,,,7629475.84,,,,,,3.12.3,2.11.0+cu128,12.8,NVIDIA GeForce RTX 5090
+Williams/webbase-1M,1000005,1000005,3105536,int64,float32,coo,sparse_logsumexp,528.8523874090364,4.086195664051818,248.271872,0.0,670.9825494752256,56.541783149895586,248.271872,0.0,3.12.3,2.11.0+cu128,12.8,NVIDIA GeForce RTX 5090
+Williams/webbase-1M,1000005,1000005,3105536,int64,float32,coo,dense.logsumexp,,,3814737.92,,,,,,3.12.3,2.11.0+cu128,12.8,NVIDIA GeForce RTX 5090
+Williams/webbase-1M,1000005,1000005,3105536,int64,float32,csr,sparse_logsumexp,581.2928689548605,3.7917376836134387,264.46694399999996,5.684341886080802e-14,767.8863043882743,87.77875094499424,264.46694399999996,5.684341886080802e-14,3.12.3,2.11.0+cu128,12.8,NVIDIA GeForce RTX 5090
+Williams/webbase-1M,1000005,1000005,3105536,int64,float32,csr,dense.logsumexp,,,3814737.92,,,,,,3.12.3,2.11.0+cu128,12.8,NVIDIA GeForce RTX 5090
+Williams/webbase-1M,1000005,1000005,3105536,int64,float64,coo,sparse_logsumexp,652.1855060371661,4.739637246101876,332.66534400000006,5.684341886080802e-14,841.7154034028387,66.40554727399672,332.66534400000006,5.684341886080802e-14,3.12.3,2.11.0+cu128,12.8,NVIDIA GeForce RTX 5090
+Williams/webbase-1M,1000005,1000005,3105536,int64,float64,coo,dense.logsumexp,,,7629475.84,,,,,,3.12.3,2.11.0+cu128,12.8,NVIDIA GeForce RTX 5090
+Williams/webbase-1M,1000005,1000005,3105536,int64,float64,csr,sparse_logsumexp,665.9346197367362,5.687582394801057,349.8306559999999,5.684341886080802e-14,927.1953889282838,107.85090915821453,349.8306559999999,5.684341886080802e-14,3.12.3,2.11.0+cu128,12.8,NVIDIA GeForce RTX 5090
+Williams/webbase-1M,1000005,1000005,3105536,int64,float64,csr,dense.logsumexp,,,7629475.84,,,,,,3.12.3,2.11.0+cu128,12.8,NVIDIA GeForce RTX 5090

torchsparsegradutils/benchmarks/results/sparse_logsumexp_suitesparse_results.csv

@@ -0,0 +1,175 @@
+#!/usr/bin/env python3
+"""
+Sparse Log-Sum-Exp Benchmark - Random Matrices
+
+Benchmarks ``sparse_logsumexp`` against a dense ``torch.logsumexp`` baseline on
+randomly generated sparse matrices of various sizes, sparsity patterns, layouts
+and dtypes. The reduction is over ``dim=1`` (one value per row); forward and
+backward time and peak memory are measured. The dense baseline materialises
+``A.to_dense()`` and is the memory ceiling the sparse path exists to avoid.
+"""
+
+import os
+import sys
+
+# Add the parent directory to sys.path to allow importing torchsparsegradutils
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
+
+import numpy as np
+import torch
+from benchmark_utils import (
+    measure_op,
+    print_benchmark_header,
+    print_result_row,
+    print_results_table_header,
+    save_benchmark_results,
+)
+from tqdm import tqdm
+
+from torchsparsegradutils import sparse_logsumexp
+from torchsparsegradutils.utils import rand_sparse
+
+REPEATS = 100
+WARMUP_RUNS = 10
+
+# Only run on CUDA
+device = torch.device("cuda")
+assert torch.cuda.is_available(), "This benchmark requires a CUDA GPU"
+
+# problem sizes: (label, N, M, nnz)
+SIZES = [
+    ("small", 2**10, 2**10, 2**12),
+    ("medium", 2**12, 2**12, 2**14),
+    ("large", 2**14, 2**14, 2**16),
+    ("xlarge", 2**16, 2**16, 2**18),
+    ("million", 2**20, 2**20, 2**22),
+]
+
+INDEX_DTYPES = [torch.int32, torch.int64]
+VALUE_DTYPES = [torch.float32, torch.float64]
+LAYOUTS = [torch.sparse_coo, torch.sparse_csr]
+
+# reduce over dim=1; measure_op takes a binary op, so B is an unused placeholder.
+ALGORITHMS = [
+    ("sparse_logsumexp", lambda A, B: sparse_logsumexp(A, dim=1)),
+    ("dense.logsumexp", lambda A, B: torch.logsumexp(A.to_dense(), dim=1)),
+]
+
+
+def run_sparse_logsumexp_benchmark():
+    """Run the sparse log-sum-exp benchmark suite."""
+
+    print_benchmark_header("Sparse Log-Sum-Exp Benchmark - Random Matrices")
+
+    records = []
+
+    for size_label, N, M, nnz in tqdm(SIZES, desc="Problem sizes"):
+        print(f"\n🔍 Testing size: {size_label} (N={N}, M={M}, nnz={nnz})")
+
+        A_shape = (N, M)
+        B = torch.zeros(1, device=device)  # unused placeholder for measure_op's binary signature
+
+        for idx_dt in tqdm(INDEX_DTYPES, desc="Index dtypes", leave=False):
+            for val_dt in tqdm(VALUE_DTYPES, desc="Value dtypes", leave=False):
+                for layout in tqdm(LAYOUTS, desc="Layouts", leave=False):
+                    layout_name = "coo" if layout == torch.sparse_coo else "csr"
+                    print(f"\n  📊 Configuration: idx_dtype={idx_dt}, val_dtype={val_dt}, layout={layout_name}")
+
+                    print_results_table_header()
+
+                    for alg_name, alg_fn in ALGORITHMS:
+                        try:
+                            print(f"    🧮 Testing {alg_name} ({layout_name})...")
+
+                            A_sparse = rand_sparse(
+                                A_shape, nnz, layout, indices_dtype=idx_dt, values_dtype=val_dt, device=device
+                            )
+
+                            (
+                                t_fwd,
+                                std_fwd,
+                                mem_fwd,
+                                std_mem_fwd,
+                                t_bwd,
+                                std_bwd,
+                                mem_bwd,
+                                std_mem_bwd,
+                            ) = measure_op(
+                                alg_fn,
+                                A_sparse,
+                                B,
+                                repeats=REPEATS,
+                                device=device,
+                                desc=f"{alg_name} ({layout_name})",
+                                warmup_runs=WARMUP_RUNS,
+                                remove_outliers=True,
+                            )
+
+                            print_result_row(
+                                f"{alg_name} ({layout_name})",
+                                (N, M),
+                                t_fwd,
+                                std_fwd,
+                                mem_fwd,
+                                std_mem_fwd,
+                                t_bwd,
+                                std_bwd,
+                                mem_bwd,
+                                std_mem_bwd,
+                            )
+
+                            records.append(
+                                {
+                                    "size": size_label,
+                                    "layout": layout_name,
+                                    "algo": alg_name,
+                                    "index_dt": str(idx_dt).split(".")[-1],
+                                    "value_dt": str(val_dt).split(".")[-1],
+                                    "N": N,
+                                    "M": M,
+                                    "nnz": nnz,
+                                    "fwd_time_us": t_fwd,
+                                    "fwd_time_std_us": std_fwd,
+                                    "fwd_mem_MB": mem_fwd,
+                                    "fwd_mem_std_MB": std_mem_fwd,
+                                    "bwd_time_us": t_bwd,
+                                    "bwd_time_std_us": std_bwd,
+                                    "bwd_mem_MB": mem_bwd,
+                                    "bwd_mem_std_MB": std_mem_bwd,
+                                }
+                            )
+
+                        except Exception as e:
+                            print(f"    ❌ {alg_name} ({layout_name}) failed: {e}")
+
+                            records.append(
+                                {
+                                    "size": size_label,
+                                    "layout": layout_name,
+                                    "algo": alg_name,
+                                    "index_dt": str(idx_dt).split(".")[-1],
+                                    "value_dt": str(val_dt).split(".")[-1],
+                                    "N": N,
+                                    "M": M,
+                                    "nnz": nnz,
+                                    "fwd_time_us": np.nan,
+                                    "fwd_time_std_us": np.nan,
+                                    "fwd_mem_MB": np.nan,
+                                    "fwd_mem_std_MB": np.nan,
+                                    "bwd_time_us": np.nan,
+                                    "bwd_time_std_us": np.nan,
+                                    "bwd_mem_MB": np.nan,
+                                    "bwd_mem_std_MB": np.nan,
+                                    "error": str(e),
+                                }
+                            )
+
+    # Save results
+    if records:
+        save_benchmark_results(records, "sparse_logsumexp_rand")
+
+    print("\n✅ Sparse log-sum-exp benchmark completed!")
+
+
+if __name__ == "__main__":
+    run_sparse_logsumexp_benchmark()