jax-ml/jax · error · ValueError

{reorder=} not valid, must be one of [1, 2, 3, 4]

Error message

{reorder=} not valid, must be one of [1, 2, 3, 4]

What it means

spsolve's reorder parameter (controlling the matrix reordering scheme passed to the GPU solver, e.g. cusolver) must be one of the integer codes 0-3. Any other value is rejected.

Source

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

  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),)

  result, _, _ = mlir.emit_python_callback(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Use an integer reorder in [0, 1, 2, 3] (typically 0 = no reorder, 1-3 = reordering schemes)
  2. If a different reordering scheme is needed, check the JAX/cusolver API version for supported codes

Example fix

// before
x = sparse.linalg.spsolve(A, b, reorder=4)
// after
x = sparse.linalg.spsolve(A, b, reorder=1)
Defensive patterns

Strategy: validation

Validate before calling

assert reorder in (0, 1, 2, 3), f'bad reorder={reorder}'

Type guard

def is_valid_reorder(r) -> bool:
    return isinstance(r, int) and 0 <= r <= 3

Prevention

When it happens

Trigger: Calling spsolve(..., reorder=k) with k outside {0,1,2,3}, e.g. passing 4 (the message text mistakenly says 1-4) or a string.

Common situations: Porting cusolver code that documents COLPERM values differently; typos or passing the parameter by keyword with a wrong constant.

Related errors


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