Add sparse_bidir_logsumexp: simultaneous row- and column-wise log-sum-exp (#89)

* refactor: generalise _scatter_logsumexp to a leading batch dim

Operate on the last axis so values/scatter_index of shape (*batch, nnz)
reduce independently into (*batch, n_groups). The 1-D path is unchanged
(dim=-1 == 0, gather(...,-1,idx) == idx-fancy-index), so sparse_logsumexp's
existing callers and their numerical fixes are untouched; enables a single
batched scatter for the upcoming sparse_bidir_logsumexp.

Suggested-by: tvercaut

* feat: add sparse_bidir_logsumexp (simultaneous row + column LSE)

Computes the column-wise (dim=0) and row-wise (dim=1) log-sum-exp of a 2-D or
batched 3-D sparse tensor in a single traversal, instead of two sparse_logsumexp
passes. Both reductions share one extraction feeding a single batched
_scatter_logsumexp via values.expand(2, nnz) (a view); its backward correctly
sums each nonzero's row+col gradient. Supports COO/CSR/CSC, include_zeros,
keepdim (tuple only), and output_layout in {tuple, padded, nested}; padded is
the native (2, G) buffer.

Suggested-by: tvercaut

* test: cover sparse_bidir_logsumexp

Parametrised over COO/CSR/CSC x value/index dtypes x include_zeros (reusing the
sparse_logsumexp fixtures/helpers): tuple equals the two-call baseline and dense
reference; padded/nested layouts agree and pad to -inf; keepdim shapes and its
rejection for non-tuple layouts; the sparse_logsumexp numerical edge cases (very-negative
values, explicit +inf) exercised on *both* axes; batched shapes; gradient parity
with the two-call baseline; and the error paths.

* bench: add sparse_bidir_logsumexp benchmarks vs two-call baseline

Random-matrix and SuiteSparse (Rothberg/cfd2, Williams/webbase-1M) benchmarks
comparing sparse_bidir_logsumexp against the two-call baseline
(sparse_logsumexp dim=0 + dim=1), mirroring the sparse_logsumexp scripts and
registered in benchmark_suite. Both ops return a single tensor (padded / cat)
so measure_op's backward exercises both reductions. Real >=1M-dim coverage per
the sparse_logsumexp review.

* docs: list sparse_bidir_logsumexp in the README key features

* docs: correct batched equivalence and soften bidir docstring claims

The 'Equivalent to' snippet used dim=0/dim=1, which raises on a batched 3-D
input (batch axis is dim=0); qualify it as 2-D and note the batched dim=1/dim=2
reductions. Also soften the unqualified 'cheaper than' (benchmarks show ~20%
higher peak memory), drop the subjective 'most pythonic', document the batched
tuple shapes, and note that the nested layout is an unbind-only prototype whose
whole-tensor ops (.sum()/.shape) are unsupported.

Suggested-by: Fable

* fix: reject hybrid sparse tensors and gate nested up front

A 3-D hybrid sparse tensor (2 sparse + 1 dense dim) has ndim==3, so it slipped
past the rank check and was misrouted as batched, then crashed on index
unpacking. Add a dense_dim() != 0 guard to both public entry points, matching
the repo's existing dense-dim validation in utils (ValueError, 'zero dense
dimensions'). Also hoist the nested PyTorch>=2.4 version gate ahead of the
reduction so it no longer pays the full forward cost before raising.

Suggested-by: Fable

* test: cover batched layouts, wide matrices, and CSC

Close the periphery gaps the review flagged: batched padded/nested now assert
values (not just shapes), so a col/row plane swap in the batched padded assembly
can't pass; a transposed wide (4x5) case exercises the row-side -inf padding that
the tall fixtures never hit; and a forward-only CSC test reaches the free-col_nnz
path that the [coo, csr] fixture set never covered.

Suggested-by: Fable

* test: add batched grad parity, duplicate-COO, and raises matchers

Cover the remaining robustness gaps: a batched gradient-parity case (the
_bidir_batched backward is a distinct path, previously only probed manually); an
uncoalesced COO with repeated coordinates (duplicates must sum before exp, so a
refactor to _values()/_indices() that double-counts would fail); and match=
patterns on the two ValueError raises so an unrelated ValueError can't satisfy
them.

Suggested-by: Fable

* bench: make index-dtype labels truthful and fail loudly

COO always stores int64 indices, so an int32 request was silently promoted and
every 'int32' COO row (and the SuiteSparse CSR rows built via to_sparse_csr)
duplicated the int64 measurement under a wrong label. Skip COO+int32 loudly,
rebuild SuiteSparse CSR with the requested index dtype, and assert the built
dtype matches the label. Also hoist matrix generation out of the timing try so a
generation failure surfaces instead of being misattributed to the op.

Suggested-by: Fable

* docs: clarify n_groups is shared across batch slices

_scatter_logsumexp's output is a rectangular (*batch, n_groups), so n_groups
is one value for every slice. Document that a caller whose slices need
different group counts passes the max and pads — the surplus groups come back
-inf and can be sliced off, which is the contract the bidirectional path
relies on with G = max(nrows, ncols).

Addresses tvercaut's review comment on PR #89.

* docs: note col_lse/row_lse are views on the padded buffer

_bidir_2d returns the scatter's native output buffer alongside two basic-slice
views into it, so the three returns cost one allocation between them rather
than three.

Addresses tvercaut's review comment on PR #89.

* perf: only pay the padded transpose when that layout is requested

_bidir_batched returned padded.permute(1, 0, 2).contiguous(), materialising a
full (b, 2, G) copy on every batched call — including output_layout="tuple" and
"nested", which discard it untouched.

It now returns the scatter's native (2, b, G) buffer with col_lse/row_lse as
views into it, matching _bidir_2d, and the public function does the transpose
only in the "padded" branch. Public behaviour is unchanged: "padded" still
returns a contiguous (b, 2, G).

The copy cannot be avoided for "padded" itself — the scatter's batch axis
carries the two directions and the slice is folded into the group index, so the
native layout is inherently (2, b*G).

Measured on a (64, 8192, 8192) batch with 500k nnz: tuple-path peak CUDA memory
drops 70.8 -> 67.2 MB, the size of the discarded buffer.

Addresses tvercaut's review comment on PR #89.

* docs: note why _logsumexp_batched folds instead of using the batch axis

_scatter_logsumexp's batch axis needs a rectangular (*batch, nnz), but the
slices have ragged nnz. Padding them to (b, max_nnz) is correct but measured
2.6x slower at 1.4x the memory with evenly-filled slices, and degrades sharply
with skew (one heavy slice gave a 24x padding factor and 27x the time).

Addresses tvercaut's review comment on PR #89.

* docs: say where _logsumexp_batched's unequal-nnz slices come from

The comment said padding "the ragged slices" was slower without saying why the
slices are ragged, three lines under a docstring stating that batched CSR/CSC
require equal nnz per slice. The two read as contradictory.

Batched COO is the layout that permits unequal nnz — one global (3, nnz) index
matrix, so slices can hold any counts — and _logsumexp_batched coalesces every
input to COO, so that case is reachable. Name the layout in the comment.

Addresses tvercaut's review comment on PR #89.

* refactor: make the padded layout direction-first at both ranks

output_layout="padded" returned (2, G) unbatched but (batch, 2, G) batched, so
out[0] was col_lse in the first case and the first batch slice in the second.
The same argument meant different things depending on input rank.

Return the scatter's native (2, batch, G) buffer instead, and document it. The
direction now leads at both ranks: out[0] and out[1] are the column and row
reductions whatever the input. This is a breaking change to the "padded" shape.

Serving the old shape required a permute + contiguous, so this also drops a copy
and every layout is now a view into the one scatter allocation. That is worth
1.6-3.7% of the call (larger G, fewer nnz favours it), and no peak memory: the
copy was transient and sat under the scatter's own high-water mark.

Addresses tvercaut's review comment on PR #89.

* docs: publish sparse_bidir_logsumexp and flag "padded" as experimental

sparse_bidir_logsumexp was exported in __all__ but never added to
docs/source/api/core.rst, so its docstring rendered nowhere.

The "padded" layout hands back _scatter_logsumexp's native output buffer
without a copy, so its axis order, G = max(rows, cols) group padding and
-inf fill are all artifacts of that buffer. Say so, and point callers who
need a stable layout at "tuple".

Addresses tvercaut's review comment on PR #89.

Suggested-by: tvercaut

* docs: align batching terminology and state the batched CSR/CSC constraint

The docstrings described inputs as "2-D" / "batched 3-D (batch, rows, cols)",
which is not the vocabulary the rest of the repo uses. Say unbatched [r, c] and
batched [b, r, c] throughout, and spell the output shapes with the same letters
((b, c), (2, b, G), G = max(r, c)).

Both functions advertise CSR/CSC support, but not that PyTorch requires every
slice of a batched compressed tensor to hold the same number of specified
elements — a ragged batch is COO-only, and the caller hits that at construction,
before reaching us. Say so where the layouts are listed.

Addresses theo-barfoot's review comments on PR #89.

Referred-by: Theo Barfoot <theo.barfoot@gmail.com>

* test: cover CSR and CSC, batched and unbatched, in both logsumexp suites

The docstrings claim COO/CSR/CSC support at both ranks, but CSC only had a
single forward-only smoke test, and every batched test hardcoded .to_sparse_coo()
— so batched CSR/CSC, which reach the reduction by a different route (converted,
as batched compressed tensors expose no per-slice index accessors), were never
exercised.

Add FWD_LAYOUTS (= SPARSE_LAYOUTS + CSC) and drive the forward tests off it,
dropping the one-off test_csc_layout it subsumes. SPARSE_LAYOUTS still backs the
gradient tests: PyTorch has no backward for CSC values.

Batched CSR/CSC cannot represent unequal nnz per slice — PyTorch raises "Expect
the same number of specified elements per batch" at construction — so the ragged
_make_batched_dense fixture is COO-only by necessity. Add
_make_batched_dense_equal_nnz (one sparsity mask shared across slices, still with
an all-zero row so empty segments are covered) as the widest input all three
layouts can express, and use it for the new batched-layout tests.

Also pin the convention behind the batched path: _nnz() is a whole-tensor count
for COO but a per-slice count for batched CSR/CSC, so batched code reading it as
a total is wrong for the compressed layouts — which is why the batched reduction
converts to COO before counting.

282 -> 510 tests, all green.

Addresses theo-barfoot's review comment on PR #89.

Referred-by: Theo Barfoot <theo.barfoot@gmail.com>

Author: aymuos15

README.md

@@ -19,6 +19,9 @@ A comprehensive collection of utility functions to work with PyTorch sparse tens
 - 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_bidir_logsumexp`: Row- and column-wise `log-sum-exp` simultaneously in a single traversal
+- Fuses the two `sparse_logsumexp(dim=0)` / `sparse_logsumexp(dim=1)` passes into one batched scatter, sharing the index extraction and autograd graph
+- Returns `tuple`, `padded`, or `nested` output layouts
 
 **Sparse Linear System Solvers**
 - `sparse_triangular_solve`: Sparse triangular solver with batch support

docs/source/api/core.rst

@@ -32,6 +32,8 @@ Sparse Reductions
 
 .. autofunction:: sparse_logsumexp
 
+.. autofunction:: sparse_bidir_logsumexp
+
 Sparse Linear Solvers
 ----------------------
 

torchsparsegradutils/__init__.py

@@ -1,5 +1,5 @@
 from .indexed_matmul import gather_mm, segment_mm
-from .sparse_logsumexp import sparse_logsumexp
+from .sparse_logsumexp import sparse_bidir_logsumexp, 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
@@ -12,4 +12,5 @@
     "sparse_generic_solve",
     "sparse_generic_lstsq",
     "sparse_logsumexp",
+    "sparse_bidir_logsumexp",
 ]

torchsparsegradutils/benchmarks/benchmark_suite.py

@@ -21,6 +21,8 @@
     "batched_sparse_mm_rand.py",
     "sparse_logsumexp_rand.py",
     "sparse_logsumexp_suitesparse.py",
+    "sparse_bidir_logsumexp_rand.py",
+    "sparse_bidir_logsumexp_suitesparse.py",
 ]
 
 

torchsparsegradutils/benchmarks/sparse_bidir_logsumexp_rand.py

@@ -0,0 +1,191 @@
+#!/usr/bin/env python3
+"""
+Sparse Bidirectional Log-Sum-Exp Benchmark - Random Matrices
+
+Benchmarks ``sparse_bidir_logsumexp`` (both row- and column-wise reductions in a
+single traversal) against the two-call baseline
+``(sparse_logsumexp(A, dim=0), sparse_logsumexp(A, dim=1))`` on randomly generated
+sparse matrices of various sizes, sparsity patterns, layouts and dtypes. Forward
+and backward time and peak memory are measured; the baseline is what the single-
+pass primitive exists to beat. Both algorithms are driven through a single tensor
+(``output_layout="padded"`` / ``torch.cat``) so ``measure_op``'s ``.sum().backward()``
+exercises both reductions.
+"""
+
+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_bidir_logsumexp, 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]
+
+# Both ops return a single tensor so measure_op's out.sum().backward() drives both
+# reductions. B is an unused placeholder for measure_op's binary signature.
+ALGORITHMS = [
+    ("sparse_bidir_logsumexp", lambda A, B: sparse_bidir_logsumexp(A, output_layout="padded")),
+    ("two_call_baseline", lambda A, B: torch.cat([sparse_logsumexp(A, dim=0), sparse_logsumexp(A, dim=1)])),
+]
+
+
+def run_sparse_bidir_logsumexp_benchmark():
+    """Run the sparse bidirectional log-sum-exp benchmark suite."""
+
+    print_benchmark_header("Sparse Bidirectional 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}")
+
+                    # COO indices are always int64 -> skip int32 rather than mislabel a duplicate.
+                    if layout == torch.sparse_coo and idx_dt != torch.int64:
+                        print(f"    ⏭  skipping {layout_name} + {idx_dt} (COO indices are always int64)")
+                        continue
+
+                    # Generate once per config, outside the try so generation errors surface here.
+                    A_sparse = rand_sparse(
+                        A_shape, nnz, layout, indices_dtype=idx_dt, values_dtype=val_dt, device=device
+                    )
+                    actual_idx_dt = (
+                        A_sparse.col_indices().dtype
+                        if layout == torch.sparse_csr
+                        else A_sparse.coalesce().indices().dtype
+                    )
+                    assert actual_idx_dt == idx_dt, f"index dtype mislabel: requested {idx_dt}, built {actual_idx_dt}"
+
+                    print_results_table_header()
+
+                    for alg_name, alg_fn in ALGORITHMS:
+                        try:
+                            print(f"    🧮 Testing {alg_name} ({layout_name})...")
+
+                            (
+                                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_bidir_logsumexp_rand")
+
+    print("\n✅ Sparse bidirectional log-sum-exp benchmark completed!")
+
+
+if __name__ == "__main__":
+    run_sparse_bidir_logsumexp_benchmark()