Skip to content

← JAX Numerical Computing step 17 of 25

Hard Primitives

Sparse Matrix-Vector Product

Jump to the task

Why this matters

Matrix-vector multiplication forms one weighted sum per row. A sparse matrix can omit work for unstored zero entries. With BCOO, sp @ vec multiplies a sparse matrix by a dense vector and returns a dense vector.

Worked mini-example

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

matrix = jnp.array([[2.0, 0.0, -1.0], [0.0, 3.0, 0.0]])
vec = jnp.array([4.0, 5.0, 6.0])
sp = BCOO.fromdense(matrix)
result = sp @ vec  # [2*4 - 1*6, 3*5] = [2, 15]

The matrix has shape (2, 3), so it consumes a length-three vector and produces two row sums.

Common pitfalls

  • Convert once when reusing a matrix. Scanning a dense matrix to construct sparse storage already costs work proportional to its size.
  • Sparse multiplication work depends on stored entries, with additional cost for indexing and the dense output. Padding and duplicate coordinates can increase that work.
  • Sparse does not guarantee faster execution. Matrix size, sparsity pattern, hardware and implementation all matter. JAX documents this module as an experimental reference implementation, not a performance-critical sparse library.
  • vec @ sp is a different operation with different shape requirements. It is supported for suitable inputs; reversing operands is not a substitute for this task.
  • Default fromdense infers its storage size from values. Build the sparse matrix before JIT, or supply a known sufficient nse for a compiled conversion.

Problem

Implement sparse_matvec(dense_matrix, vec). Convert the matrix to BCOO with its default layout, then return sp @ vec.

The finite floating-point inputs have shapes (m, n) and (n,). Return a dense vector of shape (m,), with one dot product per matrix row. An all-zero row contributes zero.

The runner calls this conversion-and-multiplication exercise directly. It demonstrates the API, not a benchmark showing that converting a dense input is faster than dense multiplication.