jax-ml/jax · error · ValueError

shape mismatch: {sparr.shape=} {a.shape=}

Error message

shape mismatch: {sparr.shape=} {a.shape=}

What it means

bcoo_extract requires the dense array's shape to equal the BCOO's shape, since it extracts values at matching indices. Any mismatch raises ValueError showing both shapes.

Source

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

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
      indices to zero-out duplicate extracted values.

  Returns:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Match shapes exactly: reshape/broadcast the dense array first (jnp.broadcast_to(dense, bcoo.shape))
2. Fix slicing/transposition of the dense operand

Example fix

// before
vals = bcoo_extract(bcoo, dense[0])  # forgot batch dim
// after
vals = bcoo_extract(bcoo, dense)

if jnp.asarray(arr).shape != sparr.shape:
    arr = jnp.broadcast_to(arr, sparr.shape)

When it happens

Trigger: bcoo_extract(bcoo, dense) where dense.shape != bcoo.shape, e.g. forgetting batch dims, transposed matrix, or leading extra axis.

Common situations: Broadcasting assumptions (extract does not broadcast); passing the unbatched dense array to a batched BCOO; shape drift after slicing.

Related errors

• 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=}
• Unsupported shape: {shape}
• batch_dims must be None or satisfy 0 < dim < n_batch. Got {b

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