{"record":{"id":"07f49cfbf743464b","repo":"jax-ml/jax","slug":"a-x-must-have-same-dtypes-were-test-output-dtyp","errorCode":null,"errorMessage":"A, X must have same dtypes (were {test_output.dtype}, {dt})","messagePattern":"A, X must have same dtypes \\(were (.+?), (.+?)\\)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/sparse/linalg.py","lineNumber":261,"sourceCode":"    assert diagnostics is not None\n    return theta[0, :], X, i, diagnostics\n  return theta[0, :], X, i\n\n\ndef _check_inputs(A, X):\n  n, k = X.shape\n  dt = X.dtype\n\n  if k == 0:\n    raise ValueError(f'must have search dim > 0, got {k}')\n\n  if k * 5 >= n:\n    raise ValueError(f'expected search dim * 5 < matrix dim (got {k * 5}, {n})')\n\n  test_output = A(jnp.zeros((n, 1), dtype=X.dtype))\n\n  if test_output.dtype != dt:\n    raise ValueError(\n        f'A, X must have same dtypes (were {test_output.dtype}, {dt})')\n\n  if test_output.shape != (n, 1):\n    s = test_output.shape\n    raise ValueError(f'A must be ({n}, {n}) matrix A, got output {s}')\n\n\ndef _mm(a, b, precision=jax.lax.Precision.HIGHEST):\n  return jax.lax.dot(a, b, precision=(precision, precision))\n\ndef _generate_diagnostics(prev_XPR, X, P, R, theta, converged, adj_resid):\n  k = X.shape[1]\n  assert X.shape == P.shape\n\n  diagdiag = lambda x: jnp.diag(jnp.diag(x))\n  abserr = lambda x: jnp.abs(x).sum() / (k ** 2)\n\n  XTX = _mm(X.T, X)","sourceCodeStart":243,"sourceCodeEnd":279,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/sparse/linalg.py#L243-L279","documentation":"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.","triggerScenarios":"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.","commonSituations":"Enabling jax_enable_x64 after building the operator; mixing float32 X with a float64 sparse matrix inside A; custom matvec with an explicit cast.","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"],"exampleFix":"// before\nA = lambda v: M @ v  # returns float64\nX = jnp.zeros((n, k), dtype=jnp.float32)\n// after\nX = X.astype(M.dtype)\n# or A = lambda v: (M @ v).astype(X.dtype)","handlingStrategy":"validation","validationCode":"out = A(jnp.zeros((X.shape[0], 1), dtype=X.dtype))\nassert out.dtype == X.dtype, f'dtype mismatch: A->{out.dtype}, X->{X.dtype}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Create the operator and X under the same x64 setting","Cast inside the matvec: lambda v: (M @ v).astype(X.dtype)"],"tags":["jax","sparse","lobpcg","dtype"],"backgroundTag":"dtype-mismatch-operator","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}