torchsparsegradutils 0.2.3__tar.gz → 0.2.5__tar.gz

{torchsparsegradutils-0.2.3/torchsparsegradutils.egg-info → torchsparsegradutils-0.2.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torchsparsegradutils
-Version: 0.2.3
+Version: 0.2.5
 Summary: A collection of utility functions to work with PyTorch sparse tensors
 Author-email: CAI4CAI research group <contact@cai4cai.uk>
 License-Expression: Apache-2.0
@@ -47,7 +47,7 @@ Dynamic: license-file
 
 # torchsparsegradutils: Sparsity-preserving gradient utility tools for PyTorch
 
-[![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)
+[![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)
 
 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.
 
@@ -61,6 +61,15 @@ 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_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
   -  Discussion reference: [PyTorch issue #87358](https://github.com/pytorch/pytorch/issues/87358)
@@ -230,15 +239,18 @@ batch_size, event_size = 2, 1000
 loc = torch.zeros(batch_size, event_size)
 
 # Example 1: LDL^T parameterization (numerically stable for precision matrices)
-# Create sparse lower triangular matrix (unit triangular, no diagonal)
+# Create sparse lower triangular matrix (strictly lower triangular, no diagonal)
+
 scale_tril = rand_sparse_tri(
     (batch_size, event_size, event_size),
-    nnz=5000,  # 5000 non-zeros for 1M parameters (0.5% sparsity)
+    nnz=5000,  # 5000 non-zeros for 1000x1000 dense matrix (99.5% sparsity)
     layout=torch.sparse_csr,
     upper=False,
-    unit_triangular=True  # Unit triangular for LDL^T
+    strict=True  # Strict triangular (exclude diagonal)
 )
 
+scale_tril.requires_grad_(True)
+
 # Diagonal component for LDL^T parameterization
 diagonal = torch.ones(batch_size, event_size) * 0.5
 
@@ -246,7 +258,7 @@ diagonal = torch.ones(batch_size, event_size) * 0.5
 dist_ldlt = SparseMultivariateNormal(
     loc=loc,
     diagonal=diagonal,
-    scale_tril=scale_tril  # Unit lower triangular
+    scale_tril=scale_tril  # Strictly lower triangular
 )
 
 # Example 2: LL^T parameterization (standard Cholesky)
@@ -255,7 +267,7 @@ scale_tril_chol = rand_sparse_tri(
     nnz=5000,
     layout=torch.sparse_csr,
     upper=False,
-    unit_triangular=False  # Include diagonal for LL^T
+    strict=False  # Include diagonal
 )
 
 # Create distribution with LL^T parameterization
@@ -270,7 +282,7 @@ precision_tril = rand_sparse_tri(
     nnz=5000,
     layout=torch.sparse_csr,
     upper=False,
-    unit_triangular=True
+    strict=True
 )
 
 precision_diagonal = torch.ones(batch_size, event_size) * 2.0
@@ -278,7 +290,7 @@ precision_diagonal = torch.ones(batch_size, event_size) * 2.0
 dist_precision = SparseMultivariateNormal(
     loc=loc,
     diagonal=precision_diagonal,
-    precision_tril=precision_tril  # Unit triangular precision factor
+    precision_tril=precision_tril  # Strictly triangular precision factor
 )
 
 # Sample with gradient support

{torchsparsegradutils-0.2.3 → torchsparsegradutils-0.2.5}/README.md

@@ -1,6 +1,6 @@
 # torchsparsegradutils: Sparsity-preserving gradient utility tools for PyTorch
 
-[![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)
+[![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)
 
 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.
 
@@ -14,6 +14,15 @@ 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_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
   -  Discussion reference: [PyTorch issue #87358](https://github.com/pytorch/pytorch/issues/87358)
@@ -183,15 +192,18 @@ batch_size, event_size = 2, 1000
 loc = torch.zeros(batch_size, event_size)
 
 # Example 1: LDL^T parameterization (numerically stable for precision matrices)
-# Create sparse lower triangular matrix (unit triangular, no diagonal)
+# Create sparse lower triangular matrix (strictly lower triangular, no diagonal)
+
 scale_tril = rand_sparse_tri(
     (batch_size, event_size, event_size),
-    nnz=5000,  # 5000 non-zeros for 1M parameters (0.5% sparsity)
+    nnz=5000,  # 5000 non-zeros for 1000x1000 dense matrix (99.5% sparsity)
     layout=torch.sparse_csr,
     upper=False,
-    unit_triangular=True  # Unit triangular for LDL^T
+    strict=True  # Strict triangular (exclude diagonal)
 )
 
+scale_tril.requires_grad_(True)
+
 # Diagonal component for LDL^T parameterization
 diagonal = torch.ones(batch_size, event_size) * 0.5
 
@@ -199,7 +211,7 @@ diagonal = torch.ones(batch_size, event_size) * 0.5
 dist_ldlt = SparseMultivariateNormal(
     loc=loc,
     diagonal=diagonal,
-    scale_tril=scale_tril  # Unit lower triangular
+    scale_tril=scale_tril  # Strictly lower triangular
 )
 
 # Example 2: LL^T parameterization (standard Cholesky)
@@ -208,7 +220,7 @@ scale_tril_chol = rand_sparse_tri(
     nnz=5000,
     layout=torch.sparse_csr,
     upper=False,
-    unit_triangular=False  # Include diagonal for LL^T
+    strict=False  # Include diagonal
 )
 
 # Create distribution with LL^T parameterization
@@ -223,7 +235,7 @@ precision_tril = rand_sparse_tri(
     nnz=5000,
     layout=torch.sparse_csr,
     upper=False,
-    unit_triangular=True
+    strict=True
 )
 
 precision_diagonal = torch.ones(batch_size, event_size) * 2.0
@@ -231,7 +243,7 @@ precision_diagonal = torch.ones(batch_size, event_size) * 2.0
 dist_precision = SparseMultivariateNormal(
     loc=loc,
     diagonal=precision_diagonal,
-    precision_tril=precision_tril  # Unit triangular precision factor
+    precision_tril=precision_tril  # Strictly triangular precision factor
 )
 
 # Sample with gradient support

{torchsparsegradutils-0.2.3 → torchsparsegradutils-0.2.5}/docs/source/conf.py

@@ -15,8 +15,8 @@ sys.path.insert(0, os.path.abspath("../../"))
 project = "torchsparsegradutils"
 copyright = "2026, CAI4CAI research group"
 author = "CAI4CAI research group"
-release = "0.2.3"
-version = "0.2.3"
+release = "0.2.5"
+version = "0.2.5"
 
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration

{torchsparsegradutils-0.2.3 → torchsparsegradutils-0.2.5}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "torchsparsegradutils"
-version = "0.2.3"
+version = "0.2.5"
 description = "A collection of utility functions to work with PyTorch sparse tensors"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -105,6 +105,12 @@ testpaths = [
 pythonpath = [
     "torchsparsegradutils/tests",
 ]
+# Fail tests when deprecated upstream APIs are exercised outside intentional compatibility paths.
+filterwarnings = [
+    'error:torch\.norm is deprecated.*:UserWarning',
+    'error:torch\.triangular_solve is deprecated.*:UserWarning',
+    'error:.*torch\.jit\.script.*is deprecated.*:DeprecationWarning',
+]
 doctest_optionflags = [
     "NORMALIZE_WHITESPACE",
     "IGNORE_EXCEPTION_DETAIL",

{torchsparsegradutils-0.2.3 → torchsparsegradutils-0.2.5}/torchsparsegradutils/__init__.py

@@ -1,4 +1,5 @@
 from .indexed_matmul import gather_mm, segment_mm
+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
@@ -10,4 +11,6 @@ __all__ = [
     "sparse_triangular_solve",
     "sparse_generic_solve",
     "sparse_generic_lstsq",
+    "sparse_logsumexp",
+    "sparse_bidir_logsumexp",
 ]

torchsparsegradutils-0.2.5/torchsparsegradutils/_compat.py

@@ -0,0 +1,48 @@
+"""Compatibility helpers for PyTorch APIs without feature-equivalent replacements."""
+
+import warnings
+
+import torch
+
+
+def linalg_solve_triangular_compat(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    *,
+    upper: bool,
+    unitriangular: bool = False,
+    transpose: bool = False,
+) -> torch.Tensor:
+    """Solve a triangular system with the appropriate dense or sparse backend.
+
+    Dense coefficient matrices use ``torch.linalg.solve_triangular``. Sparse
+    matrices continue to use the isolated legacy call because the replacement
+    API does not support the sparse layouts required by this package. See
+    https://github.com/pytorch/pytorch/issues/87358 for upstream sparse feature
+    parity tracking.
+    """
+    if A.layout == torch.strided:
+        if transpose:
+            A = A.transpose(-2, -1)
+            upper = not upper
+
+        return torch.linalg.solve_triangular(
+            A,
+            B,
+            upper=upper,
+            unitriangular=unitriangular,
+        )
+
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore",
+            message=r"torch\.triangular_solve is deprecated.*",
+            category=UserWarning,
+        )
+        return torch.triangular_solve(
+            B,
+            A,
+            upper=upper,
+            transpose=transpose,
+            unitriangular=unitriangular,
+        ).solution

{torchsparsegradutils-0.2.3 → torchsparsegradutils-0.2.5}/torchsparsegradutils/benchmarks/benchmark_suite.py

@@ -19,6 +19,10 @@ BENCHMARK_SCRIPTS = [
     "sparse_triangular_solve_rand.py",
     "sparse_triangular_solve_suitesparse.py",
     "batched_sparse_mm_rand.py",
+    "sparse_logsumexp_rand.py",
+    "sparse_logsumexp_suitesparse.py",
+    "sparse_bidir_logsumexp_rand.py",
+    "sparse_bidir_logsumexp_suitesparse.py",
 ]
 
 

torchsparsegradutils-0.2.5/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()

torchsparsegradutils-0.2.5/torchsparsegradutils/benchmarks/sparse_bidir_logsumexp_suitesparse.py

@@ -0,0 +1,191 @@
+#!/usr/bin/env python3
+"""
+Sparse Bidirectional Log-Sum-Exp Benchmark - SuiteSparse Collection
+
+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 real matrices from
+the SuiteSparse Matrix Collection (Rothberg/cfd2 ~123k x 123k, Williams/webbase-1M
+~1M x 1M). Forward and backward time and peak memory are measured. 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 (
+    load_mat_from_suitesparse_collection,
+    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
+
+REPEATS = 100
+WARMUP_RUNS = 10
+
+# Only run on CUDA
+device = torch.device("cuda")
+assert torch.cuda.is_available(), "This benchmark requires a CUDA GPU"
+
+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)])),
+]
+
+MATRICES = [
+    ("Rothberg", "cfd2"),  # 123k x 123k
+    ("Williams", "webbase-1M"),  # 1M x 1M
+]
+
+
+def run_sparse_bidir_logsumexp_suitesparse_benchmark():
+    """Run the sparse bidirectional log-sum-exp benchmark suite with SuiteSparse matrices."""
+
+    print_benchmark_header("Sparse Bidirectional Log-Sum-Exp Benchmark - SuiteSparse Collection")
+
+    records = []
+    B = torch.zeros(1, device=device)  # unused placeholder for measure_op's binary signature
+
+    for dirname, matname in tqdm(MATRICES, desc="Matrices"):
+        A_np_coo = load_mat_from_suitesparse_collection(dirname, matname)
+        matrix_name = f"{dirname}/{matname}"
+        N, M = A_np_coo.shape
+        nnz = A_np_coo.nnz
+        print(f"📊 Matrix: {matrix_name}, shape={A_np_coo.shape}, nnz={nnz}")
+
+        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🔍 Testing 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
+
+                    values = torch.tensor(A_np_coo.data, dtype=val_dt, device=device)
+                    indices = torch.tensor([A_np_coo.row, A_np_coo.col], dtype=torch.int64, device=device)
+                    A_coo = torch.sparse_coo_tensor(
+                        indices, values, A_np_coo.shape, dtype=val_dt, device=device
+                    ).coalesce()
+                    if layout == torch.sparse_csr:
+                        # to_sparse_csr() gives int64 -> rebuild with the requested dtype so the label is honest.
+                        csr = A_coo.to_sparse_csr()
+                        A_sparse = torch.sparse_csr_tensor(
+                            csr.crow_indices().to(idx_dt), csr.col_indices().to(idx_dt), csr.values(), csr.shape
+                        )
+                        assert A_sparse.col_indices().dtype == idx_dt
+                    else:
+                        A_sparse = A_coo
+
+                    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(
+                                {
+                                    "matrix": matrix_name,
+                                    "N": N,
+                                    "M": M,
+                                    "nnz": nnz,
+                                    "index_dt": str(idx_dt).split(".")[-1],
+                                    "value_dt": str(val_dt).split(".")[-1],
+                                    "layout": layout_name,
+                                    "algo": alg_name,
+                                    "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(
+                                {
+                                    "matrix": matrix_name,
+                                    "N": N,
+                                    "M": M,
+                                    "nnz": nnz,
+                                    "index_dt": str(idx_dt).split(".")[-1],
+                                    "value_dt": str(val_dt).split(".")[-1],
+                                    "layout": layout_name,
+                                    "algo": alg_name,
+                                    "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),
+                                }
+                            )
+
+    if records:
+        save_benchmark_results(records, "sparse_bidir_logsumexp_suitesparse")
+
+    print("\n✅ Sparse bidirectional log-sum-exp SuiteSparse benchmark completed!")
+
+
+if __name__ == "__main__":
+    run_sparse_bidir_logsumexp_suitesparse_benchmark()