jax-ml/jax · error · ValueError

Arrays must be one-dimensional. Got {data.shape=} {indices.s

Error message

Arrays must be one-dimensional. Got {data.shape=} {indices.shape=} {indptr.shape=} {b.shape=}

What it means

spsolve only accepts one-dimensional data, indices, indptr, and b buffers. Any of these being 2-D (e.g. a batched RHS or a matrix-shaped data buffer) fails validation.

Source

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

  # X; taking just the first m columns H(w) vstack(0, eye(m), 0) yields
  # an orthogonal extension to X.
  other = jnp.concatenate(
      [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]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use a single 1-D b; loop or vmap over columns of a multi-RHS problem
Squeeze stray size-1 dimensions from the buffers
Build the CSR without batch dims (n_batch=0)

Example fix

// before
x = sparse.linalg.spsolve(A, B)  # B shape (n, m)
// after
x = jax.vmap(lambda b: sparse.linalg.spsolve(A, b), in_axes=1, out_axes=1)(B)

Defensive patterns

Strategy: validation

Validate before calling

assert data.ndim == indices.ndim == indptr.ndim == b.ndim == 1

Prevention

vmap over the RHS columns for multi-RHS solves
Squeeze accidental extra dims before calling spsolve

When it happens

Trigger: Passing b with shape (n, m) for multi-RHS solve; a CSR object whose buffers gained an extra dimension through batching transforms (vmap without sparsify) or reshaping.

Common situations: Trying to solve for multiple right-hand sides at once; using vmap over an spsolve call in a way that leaves arrays 2-D; matrix stored with leading batch dimension.

Related errors

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