jax-ml/jax · error · ValueError

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

Error message

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

What it means

The trailing part of BCOO indices must be exactly (nse, n_sparse) with n_sparse = indices.shape[-1]; this check is effectively an internal consistency assert expressed as an error. It fires when the indices layout is malformed after the batch dims.

Source

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

  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)


#----------------------------------------------------------------------
# bcoo_todense

bcoo_todense_p = core.Primitive('bcoo_todense')

def bcoo_todense(mat: BCOO) -> Array:
  """Convert batched sparse matrix to a dense matrix.

  Args:
    mat: BCOO matrix.

  Returns:
    mat_dense: dense version of ``mat``.
  """
  return _bcoo_todense(mat.data, mat.indices, spinfo=mat._info)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Rebuild the BCOO from scratch via fromdense or round-trip through todense
2. Avoid mutating .indices/.data in place; use public bcoo_* helpers
3. Note the message has a stray '=' (format quirk); treat it as internal invariant failure

assert indices.ndim >= 2 and indices.shape[-2:] == (indices.shape[-2], indices.shape[-1])

When it happens

Trigger: Corrupt or hand-mangled indices arrays where indices.shape[n_batch:] != (nse, n_sparse); almost always reached through _validate_bcoo or extract primitives with inconsistent buffers.

Common situations: Direct buffer manipulation, JIT caching stale buffers with different nse, or a bug in code producing indices.

Related errors

• indices batch dimensions not compatible for {indices.shape=}
• 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=}
• Expected 0 < bdim <= n_batch; got {bdim=}, {n_batch=}

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