Skip to content

← JAX Numerical Computing step 18 of 25

Easy Primitives

BCOO Round-Trip

Jump to the task

Why this matters

sp.todense() rebuilds a dense array from sparse values and coordinates. This is useful for checking a small sparse result or passing it to an operation that needs a dense input.

The sparse object's shape supplies the full output dimensions, including rows and columns that contain only zeros.

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)
out = sp.todense()
# out = [[0, 4, 0], [-2, 0, 0]], shape (2, 3)

Default conversion stores each nonzero entry and reconstruction fills the other positions with zeros. For the finite inputs in this exercise, the numerical values and dtype are preserved. This is not a promise of byte-for-byte identity: for example, the sign bit of a negative zero need not survive.

Common pitfalls

  • A dense result needs storage for every position. Reconstructing a mostly empty m-by-n matrix still uses m * n elements.
  • If you explicitly choose too small an nse when constructing the sparse array, omitted values cannot be recovered.
  • Manually constructed BCOO arrays may contain duplicate coordinates. Reconstruction adds their values at that location.
  • Default fromdense infers its storage size from values and cannot do that inside ordinary JIT tracing. To compile the reconstruction, pass an already-built sparse array, or use a conversion with a known sufficient storage size.
  • Repeated conversions add work. A round-trip is a useful check, not an automatic performance optimization.

Problem

Implement sparse_roundtrip(dense). Convert the two-dimensional finite floating-point JAX array to BCOO, then return .todense().

Return the same shape, dtype and numerical values, including all-zero rows and columns. The runner calls this function directly; the default conversion does not need to be wrapped in JIT.

Use both conversions to practice the sparse API. The numeric tests compare the result with the input; they do not prove which representation your implementation used internally.