jax-ml/jax · error · NotImplementedError

BCSR from_scipy_sparse with nonzero n_dense/n_batch.

Error message

BCSR from_scipy_sparse with nonzero n_dense/n_batch.

What it means

BCSR.from_scipy_sparse only supports plain 2D CSR/COO/other scipy sparse matrices with no batch or dense (block) dimensions. Because scipy.sparse has no notion of n_batch or n_dense, passing nonzero values for these kwargs cannot be honored and raises NotImplementedError. For batched/dense-dim BCSR use BCSR.from_bcoo after constructing a BCOO.

Source

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

    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)
    indices = jnp.asarray(mat.indices).astype(index_dtype or jnp.int32)
    indptr = jnp.asarray(mat.indptr).astype(index_dtype or jnp.int32)
    return cls((data, indices, indptr), shape=mat.shape)

#--------------------------------------------------------------------
# vmappable handlers
def _bcsr_to_elt(cont, _, val, axis):
  if axis is None:
    return val
  if axis >= val.n_batch:
    raise ValueError(f"Cannot map in_axis={axis} for BCSR array with n_batch="
                     f"{val.n_batch}. in_axes for batched BCSR operations must "

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Drop n_dense/n_batch (call with defaults 0) for plain 2D scipy matrices
  2. For batched stacks, convert each scipy matrix separately and stack, or build a BCOO with the right layout and use BCSR.from_bcoo
  3. Use bcoo.bcoo_from_scipy_sparse(mat, n_batch=..., n_dense=...) if BCOO output is acceptable

Example fix

# before
m = BCSR.from_scipy_sparse(sp_mat, n_batch=1)  # NotImplementedError

# after (stack manually)
m = jax.tree.map(BCSR.from_scipy_sparse, stack_of_sp_mats)  # or per-matrix conversion
Defensive patterns

Strategy: validation

Validate before calling

assert n_dense == 0 and n_batch == 0, 'from_scipy_sparse supports no dense/batch dims'

Try / catch

try:
    m = BCSR.from_scipy_sparse(mat, n_batch=n_batch)
except NotImplementedError:
    m = BCSR.from_bcoo(bcoo.bcoo_from_scipy_sparse(mat, n_batch=n_batch))

Prevention

When it happens

Trigger: Calling BCSR.from_scipy_sparse(mat, n_batch=k) or with n_dense=k where k != 0. Signature allows the kwargs for API uniformity but only 0 is implemented.

Common situations: Copy-pasting kwargs from sparse.bcoo_from_scipy_sparse (which does support n_dense/n_batch) to from_scipy_sparse; writing format-generic loaders that pass the same options to every from_* constructor.

Related errors

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