jax-ml/jax · error · ValueError

COO must have ndim=2; got {shape=}

Error message

COO must have ndim=2; got {shape=}

What it means

The legacy jax.experimental.sparse.COO format only supports 2D matrices (row/col index buffers imply exactly two sparse dimensions). COO._empty, used by sparse.empty(format='coo') and sparse.eye(format='coo'), validates len(shape) == 2. For arbitrary-dimensional sparse arrays use BCOO.

Source

Thrown at jax/experimental/sparse/coo.py:119

    """Return a copy of the COO matrix with sorted indices.

    The matrix is sorted by row indices and column indices per row.
    If self._rows_sorted is True, this returns ``self`` without a copy.
    """
    # TODO(jakevdp): would be benefit from lowering this to cusparse sort_rows utility?
    if self._rows_sorted:
      return self
    row, col, data = lax.sort((self.row, self.col, self.data), num_keys=2)
    return self.__class__((data, row, col), shape=self.shape,
                          rows_sorted=True)

  @classmethod
  def _empty(cls, shape: Sequence[int], *, dtype: DTypeLike | None = None,
             index_dtype: DTypeLike = 'int32') -> COO:
    """Create an empty COO instance. Public method is sparse.empty()."""
    shape = tuple(shape)
    if len(shape) != 2:
      raise ValueError(f"COO must have ndim=2; got {shape=}")
    data = jnp.empty(0, dtype)
    row = col = jnp.empty(0, index_dtype)
    return cls((data, row, col), shape=shape, rows_sorted=True,
               cols_sorted=True)

  @classmethod
  def _eye(cls, N: int, M: int, k: int, *, dtype: DTypeLike | None = None,
           index_dtype: DTypeLike = 'int32') -> COO:
    if k > 0:
      diag_size = min(N, M - k)
    else:
      diag_size = min(N + k, M)

    if diag_size <= 0:
      # if k is out of range, return an empty matrix.
      return cls._empty((N, M), dtype=dtype, index_dtype=index_dtype)

    data = jnp.ones(diag_size, dtype=dtype)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use format='bcoo' (jax.experimental.sparse.BCOO) for any non-2D sparse array
2. Reshape/flatten the problem to 2D if the legacy COO API is required
3. Prefer BCOO in new code even for 2D — COO/CSR/CSC are legacy wrappers

Example fix

# before
m = sparse.empty((2, 3, 4), format='coo')  # ValueError: ndim=2 required

# after
m = sparse.empty((2, 3, 4), format='bcoo')

assert len(tuple(shape)) == 2, 'COO is 2D only; use bcoo'

def coo_shape_ok(shape) -> bool:
    return len(tuple(shape)) == 2

try:
    m = sparse.empty(shape, format='coo')
except ValueError:
    m = sparse.empty(shape, format='bcoo')

When it happens

Trigger: sparse.empty(shape, format='coo') or sparse.eye(N, format='coo') with a shape that is not length 2 — e.g. sparse.empty((2,3,4), format='coo') or sparse.eye on an ndim != 2 shape parameter.

Common situations: Older JAX code using COO being extended to batched/3D tensors; format strings chosen dynamically and hitting COO for non-matrix shapes.

Related errors

• Unsupported shape: {shape}
• matmul with object of shape {other.shape}
• 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/42f20729c974fa0f. Report an issue: GitHub.