jax-ml/jax · error · ValueError

index arrays must be integer typed; got

Error message

index arrays must be integer typed; got {indices.dtype=} {indptr.dtype=}

What it means

CSR-style index arrays address memory positions and must be integer typed. The spsolve abstract eval fires when indices or indptr has a non-integer dtype.

Solutions

Cast indices and indptr to an integer dtype (e.g. jnp.int32) before building the CSR
Prefer constructing the matrix via sparse.CSR.fromdense or from a scipy CSR whose index dtype is already integral

Example fix

// before
A = sparse.CSR((data, indices, indptr))  # indices float64
// after
A = sparse.CSR((data, indices.astype(jnp.int32), indptr.astype(jnp.int32)))

Defensive patterns

Strategy: type-guard

Validate before calling

import jnp
assert jnp.issubdtype(indices.dtype, jnp.integer), indices.dtype
assert jnp.issubdtype(indptr.dtype, jnp.integer), indptr.dtype

Type guard

def valid_index_dtype(a) -> bool:
    return jax.numpy.issubdtype(a.dtype, jax.numpy.integer)

Prevention

Always create index arrays with an explicit integer dtype
Convert scipy CSR arrays (already int32/int64) rather than hand-building from float data

When it happens

Trigger: Calling spsolve with a CSR whose indices/indptr were created as floats (e.g. from jnp.array of Python floats) or manually assembled buffers with wrong dtype.

Common situations: Manually constructing CSR buffers from external data (CSV/scipy conversion) where indices loaded as float64; slicing/computing indices with arithmetic that promotes to float.

Related errors

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

Appendix: source

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

  # X = H(w) vstack(-u vt, 0). But since H(w) is unitary its action must
  # preserve rank. Thus H(w) vstack(0, eye(n - k)) must be orthogonal to
  # 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

View on GitHub (pinned to 1e1c6a8fc0)