jax-ml/jax · error · ValueError

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

Error message

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

What it means

The legacy jax.experimental.sparse.CSR format only supports 2D matrices (n rows x m cols with indptr of length n+1). CSR._empty, backing sparse.empty(format='csr') and sparse.eye(format='csr'), validates len(shape) == 2. For higher-dimensional or batched sparse arrays use BCSR or BCOO.

Source

Thrown at jax/experimental/sparse/csr.py:88

  def dtype(self) -> np.dtype:
    return self.data.dtype

  def __init__(self, args, *, shape):
    self.data, self.indices, self.indptr = map(jnp.asarray, args)
    super().__init__(args, shape=shape)

  @classmethod
  def fromdense(cls, mat, *, nse=None, index_dtype=np.int32):
    if nse is None:
      nse = (mat != 0).sum()
    return csr_fromdense(mat, nse=nse, index_dtype=index_dtype)

  @classmethod
  def _empty(cls, shape, *, dtype=None, index_dtype='int32'):
    """Create an empty CSR instance. Public method is sparse.empty()."""
    shape = tuple(shape)
    if len(shape) != 2:
      raise ValueError(f"CSR must have ndim=2; got {shape=}")
    data = jnp.empty(0, dtype)
    indices = jnp.empty(0, index_dtype)
    indptr = jnp.zeros(shape[0] + 1, index_dtype)
    return cls((data, indices, indptr), shape=shape)

  @classmethod
  def _eye(cls, N, M, k, *, dtype=None, index_dtype='int32'):
    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)
    idx = jnp.arange(diag_size, dtype=index_dtype)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use format='bcsr' with n_batch=1 for stacks of matrices, or format='bcoo' for arbitrary layouts
2. Reshape the tensor into 2D if the legacy CSR API must be used
3. Prefer BCOO/BCSR in new code — CSR/CSC are legacy wrappers

Example fix

# before
m = sparse.empty((8, 16, 16), format='csr')  # ValueError

# after
m = sparse.empty((8, 16, 16), n_batch=1, format='bcsr')

assert len(tuple(shape)) == 2, 'CSR is 2D only; use bcsr with n_batch or bcoo'

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

try:
    m = sparse.empty(shape, format='csr')
except ValueError:
    m = sparse.empty(shape, n_batch=1, format='bcsr')

When it happens

Trigger: sparse.empty(shape, format='csr') or sparse.eye(..., format='csr') with a shape of length != 2 — e.g. sparse.empty((2,3,4), format='csr') or a 1D shape.

Common situations: Extending older 2D CSR code to batched tensors; dynamically choosing formats; migrating from scipy.sparse where ndim is always 2.

Related errors

• matmul with object of shape {other.shape}
• Unsupported shape: {shape}
• todense_transpose for {type(obj)}
• data batch dimensions not compatible for {data.shape=}, {sha
• Invalid {data.shape=} for {nse=}, {n_batch=}, {n_dense=}

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