jax-ml/jax · error · ValueError

data types do not match: {data.dtype=} {b.dtype=}

Error message

data types do not match: {data.dtype=} {b.dtype=}

What it means

Abstract evaluation for jax.experimental.sparse.linalg.spsolve requires the CSR data buffer and the right-hand side b to have identical dtypes; mixed dtypes are rejected before lowering.

Source

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

  # After some algebra, we see H(w) X = vstack(-u vt, 0)
  # Applying H(w) to both sides since H(w)^2 = I we have
  # 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))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast b to the matrix dtype before calling spsolve (or cast the matrix data)
2. Ensure both arrays are created under the same jax_enable_x64 setting

Example fix

// before
x = sparse.linalg.spsolve(A, b)  # A float64, b float32
// after
x = sparse.linalg.spsolve(A, b.astype(A.dtype))

assert b.dtype == mat.dtype, f'{b.dtype} != {mat.dtype}'
b = b.astype(mat.dtype)

When it happens

Trigger: Calling spsolve(csr_matrix, b) where the matrix was built with one dtype (e.g. float64) and b has another (float32), or vice versa.

Common situations: Building the sparse matrix from data of one precision and the RHS from another; x64 enabled for one array creation path but not the other; loading b from a file with different precision.

Related errors

• index arrays must be integer typed; got {indices.dtype=} {in
• Expected metadata dtype to be uint2, got: {meta.dtype}
• Sparse metadata format not implemented for {operand_dtype=}
• Sparse MMA unsupported for f32
• A, X must have same dtypes (were {test_output.dtype}, {dt})

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