Jax sparse is much slower than scipy on CPU, is this expected?

[refraction-ray]

My test scripts

import time
import numpy as np
import scipy.sparse

import jax
import jax.numpy as jnp
from jax.experimental import sparse

# --- Benchmark Parameters ---

N_QUBITS = 20
DIM = 2**N_QUBITS


# --- JAX Function Definitions ---

@jax.jit
def spmv(mat, vec):
    """
    JIT-compiled Sparse Matrix-Vector multiplication.
    The '@' operator dispatches to the correct JAX implementation
    based on the matrix type (BCOO or CSR).
    """
    return mat @ vec

@jax.jit
def convert_coo_to_csr(mat_coo):
    """JIT-compiled conversion from BCOO to CSR format."""
    return sparse.BCSR.from_bcoo(mat_coo)


def run_benchmark(dim):
    """
    Runs a benchmark comparing Scipy and JAX for sparse operations,
    distinguishing between JIT compilation and pure execution time for JAX.
    """
    print("=" * 70)
    print(f"Matrix Dimensions: {dim}x{dim}")
    print(f"JAX Backend: {jax.default_backend().upper()}")
    print("=" * 70)

    # --- 1. Generate Test Data ---
    print("1. Generating test data...")
    # Use numpy and scipy to generate initial data
    num_non_zero = int(dim * 10)

    # Generate COO components directly using numpy's fast random functions
    coo_data = np.random.rand(num_non_zero).astype(np.float32)
    coo_row = np.random.randint(0, dim, size=num_non_zero, dtype=np.int32)
    coo_col = np.random.randint(0, dim, size=num_non_zero, dtype=np.int32)

    # Construct the Scipy COO matrix from the generated components.
    # The constructor automatically handles duplicate (row, col) pairs by summing their data.
    scipy_mat_coo = scipy.sparse.coo_matrix((coo_data, (coo_row, coo_col)), shape=(dim, dim))
    
    numpy_vec = np.random.rand(dim).astype(np.float32)
    print("   Data generation complete.\n")

    # --- 2. Scipy Benchmark (Baseline) ---
    print("2. Benchmarking Scipy operations...")
    # Scipy operations are synchronous, so we can time them directly.

    # (a) Scipy COO -> CSR Conversion
    start_time = time.perf_counter()
    scipy_mat_csr = scipy_mat_coo.tocsr()
    end_time = time.perf_counter()
    scipy_convert_time = end_time - start_time
    print(f"   Scipy COO -> CSR Conversion Time: {scipy_convert_time:.6f} s")

    # (b) Scipy COO SpMV
    start_time = time.perf_counter()
    _ = scipy_mat_coo @ numpy_vec
    end_time = time.perf_counter()
    scipy_coo_time = end_time - start_time
    print(f"   Scipy COO SpMV Time:              {scipy_coo_time:.6f} s")

    # (c) Scipy CSR SpMV
    start_time = time.perf_counter()
    _ = scipy_mat_csr @ numpy_vec
    end_time = time.perf_counter()
    scipy_csr_time = end_time - start_time
    print(f"   Scipy CSR SpMV Time:              {scipy_csr_time:.6f} s\n")


    # --- 3. Prepare JAX Data ---
    print("3. Preparing JAX data...")
    jax_vec = jnp.array(numpy_vec)
    jax_mat_coo = sparse.BCOO(
        (jnp.array(scipy_mat_coo.data),
         jnp.stack([scipy_mat_coo.row, scipy_mat_coo.col], axis=1)),
        shape=(dim, dim)
    )
    print("   JAX data preparation complete.\n")

    # --- 4. JAX COO -> CSR Conversion Benchmark ---
    print("4. Benchmarking JAX COO -> CSR conversion...")
    # First run: includes JIT compilation time
    start_time_jax = time.perf_counter()
    jax_mat_csr = convert_coo_to_csr(jax_mat_coo)
    jax_mat_csr.data.block_until_ready()  # Ensure conversion is finished
    end_time_jax = time.perf_counter()
    jax_convert_time_1st = end_time_jax - start_time_jax
    print(f"   Conversion time (1st run, incl. compilation): {jax_convert_time_1st:.6f} s")

    # Second run: pure execution time
    start_time_jax = time.perf_counter()
    jax_mat_csr = convert_coo_to_csr(jax_mat_coo)
    jax_mat_csr.data.block_until_ready()
    end_time_jax = time.perf_counter()
    jax_convert_time_2nd = end_time_jax - start_time_jax
    print(f"   Conversion time (2nd run, pure execution):    {jax_convert_time_2nd:.6f} s\n")

    # --- 5. JAX COO SpMV Benchmark ---
    print("5. Benchmarking JAX (COO) SpMV...")
    # First run: includes JIT compilation time
    start_time_jax = time.perf_counter()
    result_jax_coo = spmv(jax_mat_coo, jax_vec)
    result_jax_coo.block_until_ready()
    end_time_jax = time.perf_counter()
    jax_coo_time_1st = end_time_jax - start_time_jax
    print(f"   COO SpMV time (1st run, incl. compilation): {jax_coo_time_1st:.6f} s")

    # Second run: pure execution time
    start_time_jax = time.perf_counter()
    result_jax_coo = spmv(jax_mat_coo, jax_vec)
    result_jax_coo.block_until_ready()
    end_time_jax = time.perf_counter()
    jax_coo_time_2nd = end_time_jax - start_time_jax
    print(f"   COO SpMV time (2nd run, pure execution):    {jax_coo_time_2nd:.6f} s\n")

    # --- 6. JAX CSR SpMV Benchmark ---
    print("6. Benchmarking JAX (CSR) SpMV...")
    # First run: JIT compiles the 'spmv' function for CSR-typed input
    start_time_jax = time.perf_counter()
    result_jax_csr = spmv(jax_mat_csr, jax_vec)
    result_jax_csr.block_until_ready()
    end_time_jax = time.perf_counter()
    jax_csr_time_1st = end_time_jax - start_time_jax
    print(f"   CSR SpMV time (1st run, incl. compilation): {jax_csr_time_1st:.6f} s")

    # Second run: pure execution time
    start_time_jax = time.perf_counter()
    result_jax_csr = spmv(jax_mat_csr, jax_vec)
    result_jax_csr.block_until_ready()
    end_time_jax = time.perf_counter()
    jax_csr_time_2nd = end_time_jax - start_time_jax
    print(f"   CSR SpMV time (2nd run, pure execution):    {jax_csr_time_2nd:.6f} s\n")

    # --- 7. Final Summary ---
    print("=" * 70)
    print("Benchmark Summary (Pure Execution Times)")
    print("-" * 70)
    print(f"{'Operation':<25} | {'Scipy Time (s)':<20} | {'JAX Time (s)':<20}")
    print("-" * 70)
    print(f"{'COO -> CSR Conversion':<25} | {scipy_convert_time:<20.6f} | {jax_convert_time_2nd:<20.6f}")
    print(f"{'COO SpMV':<25} | {scipy_coo_time:<20.6f} | {jax_coo_time_2nd:<20.6f}")
    print(f"{'CSR SpMV':<25} | {scipy_csr_time:<20.6f} | {jax_csr_time_2nd:<20.6f}")
    print("=" * 70)


if __name__ == '__main__':
    run_benchmark(dim=DIM)

My results on macstudio M2-ultra

======================================================================
Matrix Dimensions: 1048576x1048576
JAX Backend: CPU
======================================================================
1. Generating test data...
   Data generation complete.

2. Benchmarking Scipy operations...
   Scipy COO -> CSR Conversion Time: 0.234649 s
   Scipy COO SpMV Time:              0.022399 s
   Scipy CSR SpMV Time:              0.012584 s

3. Preparing JAX data...
   JAX data preparation complete.

4. Benchmarking JAX COO -> CSR conversion...
   Conversion time (1st run, incl. compilation): 4.958725 s
   Conversion time (2nd run, pure execution):    4.932628 s

5. Benchmarking JAX (COO) SpMV...
   COO SpMV time (1st run, incl. compilation): 0.340953 s
   COO SpMV time (2nd run, pure execution):    0.333047 s

6. Benchmarking JAX (CSR) SpMV...
   CSR SpMV time (1st run, incl. compilation): 0.398859 s
   CSR SpMV time (2nd run, pure execution):    0.351795 s

======================================================================
Benchmark Summary (Pure Execution Times)
----------------------------------------------------------------------
Operation                 | Scipy Time (s)       | JAX Time (s)
----------------------------------------------------------------------
COO -> CSR Conversion     | 0.234649             | 4.932628
COO SpMV                  | 0.022399             | 0.333047
CSR SpMV                  | 0.012584             | 0.351795
======================================================================

jax sparse is super efficient on GPU though, my test on 5090 below:

======================================================================
Matrix Dimensions: 1048576x1048576
JAX Backend: GPU
======================================================================
1. Generating test data...
   Data generation complete.

2. Benchmarking Scipy operations...
   Scipy COO -> CSR Conversion Time: 0.456758 s
   Scipy COO SpMV Time:              0.024571 s
   Scipy CSR SpMV Time:              0.016752 s

3. Preparing JAX data...
   JAX data preparation complete.

4. Benchmarking JAX COO -> CSR conversion...
   Conversion time (1st run, incl. compilation): 1.591573 s
   Conversion time (2nd run, pure execution):    0.012151 s

5. Benchmarking JAX (COO) SpMV...
   COO SpMV time (1st run, incl. compilation): 0.095935 s
   COO SpMV time (2nd run, pure execution):    0.000359 s

6. Benchmarking JAX (CSR) SpMV...
   CSR SpMV time (1st run, incl. compilation): 0.211560 s
   CSR SpMV time (2nd run, pure execution):    0.000400 s

======================================================================
Benchmark Summary (Pure Execution Times)
----------------------------------------------------------------------
Operation                 | Scipy Time (s)       | JAX Time (s)
----------------------------------------------------------------------
COO -> CSR Conversion     | 0.456758             | 0.012151
COO SpMV                  | 0.024571             | 0.000359
CSR SpMV                  | 0.016752             | 0.000400
======================================================================

is more than 20 times slower on cpu compared to scipy as expected?

[Aniketsy]

Just to add some context - JAX Sparse tends to be slower than SciPy on CPU because it doesn’t leverage low-level libraries like MKL or OpenBLAS, which SciPy is tightly integrated with. JAX Sparse is more focused on supporting GPU/TPU acceleration via XLA and enabling autodiff within JAX workflows. For performance on CPU, SciPy is usually the better choice, while JAX Sparse is more suitable when working with accelerators or when differentiability is needed. This performance difference is a known and expected limitation at the current stage of development.

[jakevdp]

jax.experimental.sparse is not all that performant on accelerators either (actually it's quite a bit worse than on CPU): it's really just a set of experimental reference implementations to explore how JAX transformations might interact with sparse operations; it's not something that we've spent time optimizing for performance on any backend.

[Aniketsy]

Thanks for the clarification .

[jakevdp]

Yes, it's expected that jax.experimental.sparse will not be particularly performant – this is called out at the top of the module documentation:

Note

The methods in jax.experimental.sparse are experimental reference implementations, and not recommended for use in performance-critical applications.

But the fact that you asked this question means that the docs didn't get that message across very effectively: if you have ideas of how to better communicate this, please let us know!

[refraction-ray]

Thanks for the reply. The sparse implementation is still super more performant than scipy on cpu as benchmarked above :) Besides, is these any alternative sparse implementation that I can use which could be more performant, with GPU/JIT/AD/VMAP support within jax ecosystem (or at least support GPU+JIT+AD)?