jax-ml/jax · error · ValueError

Invalid = for

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.

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
• batch_dims must be None or satisfy 0 < dim < n_batch. Got
• out of range for indices with
• bcoo_dynamic_slice: indices must have size mat.ndim=
• bcoo_extract with unequal batch dimensions.

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

Appendix: 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)