jax-ml/jax · error · ValueError

A, X must have same dtypes (were {test_output.dtype}, {dt})

Error message

A, X must have same dtypes (were {test_output.dtype}, {dt})

What it means

LOBPCG validates that applying the linear operator A to a zero vector preserves dtype; if A returns a different dtype than X (e.g. float32 vs float64), mixed-precision iteration is unsupported and the check fails.

Source

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

    assert diagnostics is not None
    return theta[0, :], X, i, diagnostics
  return theta[0, :], X, i


def _check_inputs(A, X):
  n, k = X.shape
  dt = X.dtype

  if k == 0:
    raise ValueError(f'must have search dim > 0, got {k}')

  if k * 5 >= n:
    raise ValueError(f'expected search dim * 5 < matrix dim (got {k * 5}, {n})')

  test_output = A(jnp.zeros((n, 1), dtype=X.dtype))

  if test_output.dtype != dt:
    raise ValueError(
        f'A, X must have same dtypes (were {test_output.dtype}, {dt})')

  if test_output.shape != (n, 1):
    s = test_output.shape
    raise ValueError(f'A must be ({n}, {n}) matrix A, got output {s}')


def _mm(a, b, precision=jax.lax.Precision.HIGHEST):
  return jax.lax.dot(a, b, precision=(precision, precision))

def _generate_diagnostics(prev_XPR, X, P, R, theta, converged, adj_resid):
  k = X.shape[1]
  assert X.shape == P.shape

  diagdiag = lambda x: jnp.diag(jnp.diag(x))
  abserr = lambda x: jnp.abs(x).sum() / (k ** 2)

  XTX = _mm(X.T, X)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Cast X to match A's output dtype (or cast inside A's matvec to X.dtype)
Build the sparse matrix and X with the same dtype
Set jax.config.update('jax_enable_x64', ...) consistently before creating both

Example fix

// before
A = lambda v: M @ v  # returns float64
X = jnp.zeros((n, k), dtype=jnp.float32)
// after
X = X.astype(M.dtype)
# or A = lambda v: (M @ v).astype(X.dtype)

Defensive patterns

Strategy: validation

Validate before calling

out = A(jnp.zeros((X.shape[0], 1), dtype=X.dtype))
assert out.dtype == X.dtype, f'dtype mismatch: A->{out.dtype}, X->{X.dtype}'

Prevention

Create the operator and X under the same x64 setting
Cast inside the matvec: lambda v: (M @ v).astype(X.dtype)

When it happens

Trigger: Calling lobpcg_standard where A is a matvec function that upcasts/downcasts (e.g. sparse matmul promoting to float64) while X is float32, or vice versa.

Common situations: Enabling jax_enable_x64 after building the operator; mixing float32 X with a float64 sparse matrix inside A; custom matvec with an explicit cast.

Related errors

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