jax-ml/jax · error · NotImplementedError

BSCR.from_bcoo requires n_sparse=2; got {arr.n_sparse=}

Error message

BSCR.from_bcoo requires n_sparse=2; got {arr.n_sparse=}

What it means

BCSR.from_bcoo converts a BCOO array into the BCSR format, which requires exactly 2 sparse dimensions (one indptr/row dim and one column-index dim). If the BCOO array has n_sparse != 2, the conversion is structurally impossible and a NotImplementedError is raised. Convert to a 2-sparse-dim BCOO first, or keep BCOO.

Source

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

  @classmethod
  def fromdense(cls, mat, *, nse=None, index_dtype=np.int32, n_dense=0,
                n_batch=0):
    """Create a BCSR array from a (dense) :class:`Array`."""
    return bcsr_fromdense(mat, nse=nse, index_dtype=index_dtype,
                          n_dense=n_dense, n_batch=n_batch)

  def todense(self):
    """Create a dense version of the array."""
    return bcsr_todense(self)

  def to_bcoo(self) -> bcoo.BCOO:
    coo_indices = _bcsr_to_bcoo(self.indices, self.indptr, shape=self.shape)
    return bcoo.BCOO((self.data, coo_indices), shape=self.shape)

  @classmethod
  def from_bcoo(cls, arr: bcoo.BCOO) -> BCSR:
    if arr.n_sparse != 2:
      raise NotImplementedError(f"BSCR.from_bcoo requires n_sparse=2; got {arr.n_sparse=}")
    if not arr.indices_sorted:
      arr = arr.sort_indices()
    indices, indptr = _bcoo_to_bcsr(
        arr.indices, shape=arr.shape, index_dtype=arr.indices.dtype
    )
    return cls((arr.data, indices, indptr), shape=arr.shape)

  @classmethod
  def from_scipy_sparse(cls, mat, *, index_dtype=None, n_dense=0, n_batch=0):
    """Create a BCSR array from a :mod:`scipy.sparse` array."""
    if n_dense != 0 or n_batch != 0:
      raise NotImplementedError("BCSR from_scipy_sparse with nonzero n_dense/n_batch.")

    if mat.ndim != 2:
      raise ValueError(f"BCSR from_scipy_sparse requires 2D array; {mat.ndim}D is given.")

    mat = mat.tocsr()
    data = jnp.asarray(mat.data)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Construct the source BCOO with n_batch/n_dense set so exactly 2 dims remain sparse
If the extra dims are batch dims, rebuild the BCOO with n_batch=... or use bcoo.reshape_to_batched to move leading dims into batch dims before converting
Keep the array as BCOO if you truly need n_sparse != 2
Wrap conversions in try/except NotImplementedError and fall back to BCOO paths

Example fix

# before
x = bcoo.bcoo_fromdense(dense_3d)  # n_sparse=3
y = BCSR.from_bcoo(x)  # NotImplementedError

# after
x = bcoo.reshape_to_batched(bcoo.bcoo_fromdense(dense_3d), 1)  # 1 batch + 2 sparse
y = BCSR.from_bcoo(x)

Defensive patterns

Strategy: validation

Validate before calling

assert arr.n_sparse == 2, f'n_sparse={arr.n_sparse}; move leading dims to n_batch or use BCOO'

Type guard

def bcoo_is_bcsr_convertible(arr) -> bool:
    return arr.n_sparse == 2

Try / catch

try:
    b = BCSR.from_bcoo(arr)
except NotImplementedError:
    b = arr  # keep as BCOO

Prevention

Construct BCOOs destined for BCSR with explicit n_batch/n_dense from the start
Check arr.n_sparse before converting
Use bcoo.reshape_to_batched to fix layout before conversion

When it happens

Trigger: Calling BCSR.from_bcoo(bcoo_array) where bcoo_array.n_sparse (len(shape) - n_dense - n_batch) is not 2 — e.g. a 3-sparse-dim BCOO, or a 1D-sparse BCOO vector. Indirectly hit via BCSR operations (todense, matvec, eliminate_zeros, sum_duplicates, broadcast_in_dim, concatenate) that internally convert from BCOO with the wrong layout.

Common situations: Building a BCOO with all dimensions sparse (the default) and then feeding it to BCSR; batched pipelines where batch dims were not declared via n_batch in the BCOO.

Related errors

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