jax-ml/jax · error · ValueError

Invalid CSR buffer sizes: {data.shape=} {indices.shape=} {in

Error message

Invalid CSR buffer sizes: {data.shape=} {indices.shape=} {indptr.shape=}

What it means

spsolve validates CSR buffer consistency: indptr must have exactly b.size + 1 entries (one per row plus sentinel) and data.shape must equal indices.shape. Violations indicate a malformed CSR structure.

Source

Thrown at jax/experimental/sparse/linalg.py:531

      [jnp.eye(m, dtype=X.dtype),
       jnp.zeros((n - k - m, m), dtype=X.dtype)], axis=0)
  w = _mm(y, vt.T * ((2 * (1 + s)) ** (-1/2))[jnp.newaxis, :])
  h = -2 * jnp.linalg.multi_dot(
      [w, w[k:, :].T, other], precision=jax.lax.Precision.HIGHEST)
  return h.at[k:].add(other)


# Sparse direct solve via QR factorization
def _spsolve_abstract_eval(data, indices, indptr, b, *, tol, reorder):
  if data.dtype != b.dtype:
    raise ValueError(f"data types do not match: {data.dtype=} {b.dtype=}")
  if not (jnp.issubdtype(indices.dtype, jnp.integer) and jnp.issubdtype(indptr.dtype, jnp.integer)):
    raise ValueError(f"index arrays must be integer typed; got {indices.dtype=} {indptr.dtype=}")
  if not data.ndim == indices.ndim == indptr.ndim == b.ndim == 1:
    raise ValueError("Arrays must be one-dimensional. "
                     f"Got {data.shape=} {indices.shape=} {indptr.shape=} {b.shape=}")
  if indptr.size != b.size + 1 or  data.shape != indices.shape:
    raise ValueError(f"Invalid CSR buffer sizes: {data.shape=} {indices.shape=} {indptr.shape=}")
  if reorder not in [0, 1, 2, 3]:
    raise ValueError(f"{reorder=} not valid, must be one of [1, 2, 3, 4]")
  tol = float(tol)
  return b


def _spsolve_gpu_lowering(ctx, data, indices, indptr, b, *, tol, reorder):
  return ffi.ffi_lowering("cusolver_csrlsvqr_ffi")(
      ctx, data, indices, indptr, b, tol=np.float64(tol),
      reorder=np.int32(reorder))

def _spsolve_cpu_lowering(ctx, data, indices, indptr, b, tol, reorder):
  del tol, reorder
  args = [data, indices, indptr, b]

  def _callback(data, indices, indptr, b, **kwargs):
    A = scipy.sparse.csr_matrix((data, indices, indptr), shape=(b.size, b.size))
    return (scipy.sparse.linalg.spsolve(A, b).astype(b.dtype),)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Rebuild the CSR from a dense array or a valid scipy CSR to guarantee consistent buffers
  2. Check indptr has length nrows+1 and len(data) == len(indices) == indptr[-1]
  3. Use sparse.CSR / sparse.bcoo_tocsr helpers instead of manual buffers

Example fix

// before
A = sparse.CSR((data, indices, indptr))  # indptr from another matrix
// after
A = sparse.CSR.fromdense(M_dense)  # or sparse.bcoo_tocsr(bcoo)
Defensive patterns

Strategy: validation

Validate before calling

assert indptr.size == b.size + 1, f'{indptr.size} != {b.size + 1}'
assert data.shape == indices.shape

Type guard

def is_valid_csr(data, indices, indptr, b) -> bool:
    return (indptr.size == b.size + 1
            and data.shape == indices.shape
            and int(indptr[-1]) == data.size)

Prevention

When it happens

Trigger: Passing indptr of the wrong length (e.g. built for a different number of rows) or data/indices arrays of differing lengths (nse mismatch).

Common situations: Hand-assembling CSR buffers from mismatched arrays; slicing the matrix rows without adjusting indptr; converting from scipy with an off-by-one in indptr.

Related errors


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