jax-ml/jax · error · ValueError

Invalid {data.shape=} for {nse=}, {n_batch=}, {n_dense=}

Error message

Invalid {data.shape=} for {nse=}, {n_batch=}, {n_dense=}

What it means

Beyond the batch dims, a BCOO's data must have shape (nse,) + dense dims, where nse comes from indices.shape[-2]. If it doesn't, _validate_bcoo raises this ValueError.

Source

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

  n_sparse: int
  n_dense: int
  nse: int

class Buffer(Protocol):
  @property
  def shape(self) -> Shape: ...
  @property
  def dtype(self) -> Any: ...


def _validate_bcoo(data: Buffer, indices: Buffer, shape: Sequence[int]) -> BCOOProperties:
  props = _validate_bcoo_indices(indices, shape)
  n_batch, n_sparse, n_dense, nse = props
  shape = tuple(shape)
  if any(s1 not in (1, s2) for s1, s2 in safe_zip(data.shape[:n_batch], shape[:n_batch])):
    raise ValueError(f"data batch dimensions not compatible for {data.shape=}, {shape=}")
  if data.shape[n_batch:] != (nse,) + shape[n_batch + n_sparse:]:
    raise ValueError(f"Invalid {data.shape=} for {nse=}, {n_batch=}, {n_dense=}")
  return props


def _validate_bcoo_indices(indices: Buffer, shape: Sequence[int]) -> BCOOProperties:
  assert jnp.issubdtype(indices.dtype, jnp.integer)
  shape = tuple(shape)
  nse, n_sparse = indices.shape[-2:]
  n_batch = len(indices.shape) - 2
  n_dense = len(shape) - n_batch - n_sparse
  assert n_dense >= 0
  if any(s1 not in (1, s2) for s1, s2 in safe_zip(indices.shape[:n_batch], shape[:n_batch])):
    raise ValueError(f"indices batch dimensions not compatible for {indices.shape=}, {shape=}")
  if indices.shape[n_batch:] != (nse, n_sparse):
    raise ValueError(f"Invalid ={indices.shape=} for {nse=}, {n_batch=}, {n_dense=}")
  return BCOOProperties(n_batch=n_batch, n_sparse=n_sparse, n_dense=n_dense, nse=nse)


#----------------------------------------------------------------------

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make data length equal indices.shape[-2] (nse)
2. Regenerate the pair via BCOO.fromdense(...) or jax.experimental.sparse.bcoo.bcoo_todense round-trip
3. Use bcoo_update_layout to fix layout instead of hand-editing buffers

Example fix

// before
m = BCOO((data[:5], indices), shape=shape)  # indices has nse=8
// after
m = BCOO((data[:5], indices[:, :5]), shape=shape)

nse = indices.shape[-2]
assert data.shape[n_batch:] == (nse,) + shape[n_batch+n_sparse:], 'data/nse mismatch'

When it happens

Trigger: Building BCOO where data.shape[n_batch:] != (nse,) + shape[n_batch+n_sparse:], e.g. data has length 10 but indices imply nse=8, or dense dims missing.

Common situations: Trimming data without trimming indices; wrong nse after eliminate_zeros/sum_duplicates done manually; forgetting trailing dense dims for dense-sparse hybrids.

Related errors

• data batch dimensions not compatible for {data.shape=}, {sha
• indices batch dimensions not compatible for {indices.shape=}
• shape mismatch: {sparr.shape=} {a.shape=}
• got {nse=}, expected to be between 0 and {sparse_size}
• Unsupported shape: {shape}

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