jax-ml/jax · error · ValueError

BCSR sparse.empty: must have 2 sparse dimensions.

Error message

BCSR sparse.empty: must have 2 sparse dimensions.

What it means

jax.experimental.sparse.BCSR (batched CSR) supports exactly 2 sparse dimensions by construction: the CSR format encodes one row-pointer (indptr) dimension and one column-index dimension. sparse.empty()/BCSR._empty computes n_sparse = len(shape) - n_dense - n_batch and rejects anything other than 2. Use BCOO if you need a different number of sparse dimensions.

Source

Thrown at jax/experimental/sparse/bcsr.py:927

  @classmethod
  def tree_unflatten(cls, aux_data, children):
    obj = object.__new__(cls)
    obj.data, obj.indices, obj.indptr = children
    if aux_data.keys() != {'shape', 'indices_sorted', 'unique_indices'}:
      raise ValueError(f"BCSR.tree_unflatten: invalid {aux_data=}")
    obj.__dict__.update(**aux_data)
    return obj

  @classmethod
  def _empty(cls, shape, *, dtype=None, index_dtype='int32', n_dense=0,
             n_batch=0, nse=0):
    """Create an empty BCSR instance. Public method is sparse.empty()."""
    shape = tuple(shape)
    if n_dense < 0 or n_batch < 0 or nse < 0:
      raise ValueError(f"Invalid inputs: {shape=}, {n_dense=}, {n_batch=}, {nse=}")
    n_sparse = len(shape) - n_dense - n_batch
    if n_sparse != 2:
      raise ValueError("BCSR sparse.empty: must have 2 sparse dimensions.")
    batch_shape, sparse_shape, dense_shape = split_list(shape,
                                                        [n_batch, n_sparse])
    data = jnp.zeros((*batch_shape, nse, *dense_shape), dtype)
    indices = jnp.full((*batch_shape, nse), jnp.array(sparse_shape[1]),
                       index_dtype)
    indptr = jnp.zeros((*batch_shape, sparse_shape[0] + 1), index_dtype)
    return cls((data, indices, indptr), shape=shape)

  def sum_duplicates(self, nse: int | None = None, remove_zeros: bool = True) -> BCSR:
    """Return a copy of the array with duplicate indices summed.

    Additionally, this operation will result in explicit zero entries removed, and
    indices being sorted in lexicographic order.

    Because the size of the resulting representation depends on the values in the
    arrays, this operation is not compatible with JIT or other transforms. To use
    ``sum_duplicates`` in such cases, you may pass a value to `nse` to specify the
    desired size of the output representation.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. If leading dims are batch dimensions, pass n_batch equal to the number of batch dims so the remaining sparse part is 2D (e.g. sparse.empty((B, M, N), n_batch=1, format='bcsr'))
  2. If trailing dims are dense, pass n_dense so the sparse part is 2D
  3. If the shape genuinely isn't 2D-sparse (e.g. 3 fully-sparse dims or 1D), use sparse.empty(..., format='bcoo') instead
  4. Catch ValueError and fall back to BCOO when format choice is dynamic

Example fix

# before
m = sparse.empty((8, 16, 16), format='bcsr')  # ValueError: must have 2 sparse dimensions

# after (first dim is a batch dim)
m = sparse.empty((8, 16, 16), n_batch=1, format='bcsr')
Defensive patterns

Strategy: validation

Validate before calling

from jax.experimental import sparse
n_sparse = len(shape) - n_dense - n_batch
assert n_sparse == 2, f'BCSR needs 2 sparse dims, got {n_sparse}; use bcoo'

Type guard

def is_bcsr_compatible(shape, n_dense=0, n_batch=0) -> bool:
    return len(tuple(shape)) - n_dense - n_batch == 2

Try / catch

try:
    m = sparse.empty(shape, n_batch=n_batch, format='bcsr')
except ValueError:
    m = sparse.empty(shape, format='bcoo')

Prevention

When it happens

Trigger: Calling sparse.empty(shape, format='bcsr') (or BCSR-related empty paths) where len(shape) minus n_dense and n_batch is not 2 — e.g. a 3D shape with n_batch=0, n_dense=0, or a 1D/4D shape, or passing n_dense/n_batch values that leave != 2 sparse dims.

Common situations: Migrating code from BCOO (which supports arbitrary n_sparse) to BCSR; passing a batched shape while forgetting to set n_batch so the batch dim is counted as a sparse dim; generically dispatching empty() over many formats with a fixed shape.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/244b2cef78abfd79. Report an issue: GitHub.