jax-ml/jax · error · TypeError

bcoo_slice: input should be BCOO array, got type(mat)={type(

Error message

bcoo_slice: input should be BCOO array, got type(mat)={type(mat)}

What it means

bcoo_slice requires the input to be a jax.experimental.sparse.BCOO array; any other type (dense jnp.ndarray, JAX Sparse (COO/CSR/CSC), numpy array) is rejected with a TypeError before any slicing logic runs. This mirrors lax.slice semantics but for the sparse BCOO representation. Convert your array to BCOO first.

Source

Thrown at jax/experimental/sparse/bcoo.py:1976

def bcoo_slice(mat: BCOO, *, start_indices: Sequence[int], limit_indices: Sequence[int],
               strides: Sequence[int] | None = None) -> BCOO:
  """Sparse implementation of :func:`jax.lax.slice`.

  Args:
    mat: BCOO array to be reshaped.
    start_indices: sequence of integers of length `mat.ndim` specifying the starting
      indices of each slice.
    limit_indices: sequence of integers of length `mat.ndim` specifying the ending
      indices of each slice
    strides: (not implemented) sequence of integers of length `mat.ndim` specifying
      the stride for each slice

  Returns:
    out: BCOO array containing the slice.
  """
  if not isinstance(mat, BCOO):
    raise TypeError(f"bcoo_slice: input should be BCOO array, got type(mat)={type(mat)}")
  start_indices = [operator.index(i) for i in start_indices]
  limit_indices = [operator.index(i) for i in limit_indices]
  if strides is not None:
    strides = [operator.index(i) for i in strides]
  else:
    strides = [1] * mat.ndim
  if len(start_indices) != len(limit_indices) != len(strides) != mat.ndim:
    raise ValueError(f"bcoo_slice: indices must have size mat.ndim={mat.ndim}")
  if len(strides) != mat.ndim:
    raise ValueError(f"len(strides) = {len(strides)}; expected {mat.ndim}")
  if any(s <= 0 for s in strides):
    raise ValueError(f"strides must be a sequence of positive integers; got {strides}")

  if not all(0 <= start <= end <= size
             for start, end, size in safe_zip(start_indices, limit_indices, mat.shape)):
    raise ValueError(f"bcoo_slice: invalid indices. Got {start_indices=}, "
                     f"{limit_indices=} and shape={mat.shape}")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Convert the operand: mat = jax.experimental.sparse.BCOO.fromdense(mat)
  2. If coming from another sparse format, convert via BCOO: sparse.BCOO.from_scipy_sparse(m) or reconstruct with BCOO((data, indices), shape=...)
  3. ,If you intended dense slicing, use lax.slice / mat[start:stop] directly instead of bcoo_slice

Example fix

// before
out = sparse.bcoo_slice(dense_mat, (0,), (4,))
// after
out = sparse.bcoo_slice(sparse.BCOO.fromdense(dense_mat), (0,), (4,))
Defensive patterns

Strategy: type-guard

Validate before calling

from jax.experimental import sparse
if not isinstance(mat, sparse.BCOO):
    mat = sparse.BCOO.fromdense(mat) if hasattr(mat, 'ndim') else mat

Type guard

import jax.experimental.sparse as sparse
from jax.experimental.sparse import BCOO
def is_bcoo(x) -> bool:
    return isinstance(x, BCOO)

Prevention

When it happens

Trigger: Calling jax.experimental.sparse.bcoo_slice(mat, start_indices, limit_indices) (or sparse.BCOO-safe code paths) with mat being a dense jnp.ndarray, np.ndarray, or a jax_sparse COO/CSR object instead of sparse.BCOO.

Common situations: Mixing jax_sparse (older external library) with jax.experimental.sparse; passing a dense array returned by a previous dense computation; refactoring code from jnp slicing to sparse slicing without converting the operand.

Related errors


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