jax-ml/jax · error · TypeError

First argument to bcoo_extract should be a BCOO array. Got {

Error message

First argument to bcoo_extract should be a BCOO array. Got {type(sparr)=}

What it means

jax.experimental.sparse.bcoo_extract requires its first argument to be a BCOO instance; it extracts the dense array's values at the BCOO's stored indices. Passing anything else (dense array, CSR, BCSR) raises TypeError.

Source

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

bcoo_extract_p = core.Primitive('bcoo_extract')


def bcoo_extract(sparr: BCOO, arr: ArrayLike, *, assume_unique: bool | None = None) -> BCOO:
  """Extract values from a dense array according to the sparse array's indices.

  Args:
    sparr : BCOO array whose indices will be used for the output.
    arr : ArrayLike with shape equal to self.shape
    assume_unique : bool, defaults to sparr.unique_indices
      If True, extract values for every index, even if index contains duplicates.
      If False, duplicate indices will have their values summed and returned in
      the position of the first index.

  Returns:
    extracted : a BCOO array with the same sparsity pattern as self.
  """
  if not isinstance(sparr, BCOO):
    raise TypeError(f"First argument to bcoo_extract should be a BCOO array. Got {type(sparr)=}")
  a = jnp.asarray(arr)
  if a.shape != sparr.shape:
    raise ValueError(f"shape mismatch: {sparr.shape=} {a.shape=}")
  if assume_unique is None:
    assume_unique = sparr.unique_indices
  data = _bcoo_extract(sparr.indices, a, assume_unique=assume_unique)
  return BCOO((data, sparr.indices), **sparr._info._asdict())


def _bcoo_extract(indices: Array, arr: Array, *, assume_unique=True) -> Array:
  """Extract BCOO data values from a dense array at given BCOO indices.

  Args:
    indices: An ndarray; see BCOO indices.
    arr: A dense array.
    assume_unique: bool, default=True
      If True, then indices will be assumed unique and a value will be extracted
      from arr for each index. Otherwise, extra work will be done to de-duplicate

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass a BCOO: convert with sparr.tobcoo() if needed
2. For raw indices, use the internal _bcoo_extract(indices, arr) instead of the public wrapper
3. For BCSR use bcsr_extract

Example fix

// before
vals = bcoo_extract(csr_mat, dense_arr)
// after
vals = bcoo_extract(csr_mat.tobcoo(), dense_arr)

from jax.experimental.sparse import BCOO
if not isinstance(sparr, BCOO):
    sparr = sparr.tobcoo() if hasattr(sparr, 'tobcoo') else BCOO.fromdense(jnp.asarray(sparr))

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

try:
    bcoo_extract(sparr, arr)
except TypeError:
    bcoo_extract(sparr.tobcoo(), arr)

When it happens

Trigger: Calling bcoo_extract(csr_mat, dense_arr), bcoo_extract(jnp.array(...), ...), or feeding a BCSR where a BCOO is required.

Common situations: Assuming bcoo_extract works for any sparse format; passing a raw indices array; format confusion when porting between CSR and BCOO code paths.

Related errors

• transpose permutation must be a tuple/list/ndarray, got {typ
• batch_dims must be None or satisfy 0 < dim < n_batch. Got {b
• data batch dimensions not compatible for {data.shape=}, {sha
• Invalid {data.shape=} for {nse=}, {n_batch=}, {n_dense=}
• indices batch dimensions not compatible for {indices.shape=}

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