Skip to content

← JAX Numerical Computing step 16 of 25

Medium Primitives

BCOO from Dense

Jump to the task

Why this matters

A sparse matrix avoids storing a value at every position. In JAX's experimental BCOO format, stored values are paired with coordinates; positions with no entry are zero.

For a two-dimensional matrix with the default layout, sp.data contains the values and sp.indices contains their row-column coordinates.

Worked mini-example

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

dense = jnp.array([[0.0, 4.0, 0.0], [-2.0, 0.0, 0.0]])
sp = BCOO.fromdense(dense)
# sp.data = [4, -2]
# sp.indices = [[0, 1], [1, 0]]
# sp.shape = (2, 3), sp.nse = 2

There are two nonzero values, including the negative value. fromdense without an explicit storage size finds those entries.

Common pitfalls

  • nse means the number of stored entries. For the default conversion above, it equals the nonzero count. For a general BCOO array it may include padding, explicit zeros or duplicate coordinates, so it is not always a nonzero count.
  • Setting nse explicitly reserves that many entries. Too few drops nonzero entries; extra slots are padded.
  • Inferring nse depends on input values. Ordinary JIT tracing cannot infer this variable storage shape from a dynamic dense input. A jitted conversion needs a known nse large enough for the input, or you can build the sparse array before JIT.
  • Sparse storage also needs coordinates. A small or nearly full matrix is not guaranteed to save memory.
  • BCOO is experimental and supports only some JAX operations.

Problem

Implement sparse_nnz(dense) for a two-dimensional finite floating-point JAX array. Create a BCOO array using the default BCOO.fromdense(dense), then return sp.nse as a jnp.float32 scalar.

Count entries unequal to zero, not their sum or the number of positive entries. An all-zero array returns 0.0.

The runner calls this exercise directly. Do not wrap the default value-dependent conversion in jax.jit. Supplying a fixed capacity and returning that capacity would not answer the task.