jax-ml/jax · error · ValueError

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

Error message

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

What it means

spsolve requires the CSR index arrays (indices and indptr) to be of integer dtype (e.g. int32/int64). Float or other dtypes for the index buffers are rejected during abstract evaluation.

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)

Solutions

1. Cast indices and indptr to an integer dtype (e.g. jnp.int32) before building the CSR
2. 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)))

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

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

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

• data types do not match: {data.dtype=} {b.dtype=}
• index arguments to dynamic_update_slice must be integers of
• Expected metadata dtype to be uint2, got: {meta.dtype}
• Sparse metadata format not implemented for {operand_dtype=}
• Sparse MMA unsupported for f32

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