{"record":{"id":"4e9b1faa81dee4a5","repo":"jax-ml/jax","slug":"qdwh-implementation-is-only-supported-on-tpu","errorCode":null,"errorMessage":"QDWH implementation is only supported on TPU","messagePattern":"QDWH implementation is only supported on TPU","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/linalg.py","lineNumber":1293,"sourceCode":"  return (n, d), (d,)\n\ndef _eigh_dtype_rule(dtype, **_):\n  return dtype, lax._complex_basetype(dtype)\n\ndef _eigh_cpu_gpu_lowering(\n    ctx, operand, *, lower, sort_eigenvalues, subset_by_index, algorithm,\n    target_name_prefix: str\n):\n  del sort_eigenvalues  # The CPU/GPU implementations always sort.\n  operand_aval, = ctx.avals_in\n  v_aval, w_aval = ctx.avals_out\n  n = operand_aval.shape[-1]\n  if not (subset_by_index is None or subset_by_index == (0, n)):\n    raise NotImplementedError(\"subset_by_index not supported on CPU and GPU\")\n  batch_dims = operand_aval.shape[:-2]\n\n  if algorithm == EighImplementation.QDWH:\n    raise NotImplementedError(\"QDWH implementation is only supported on TPU\")\n  if algorithm == EighImplementation.JACOBI and target_name_prefix == \"cpu\":\n    raise NotImplementedError(\"Jacobi implementation is not supported on CPU\")\n\n  if target_name_prefix == \"cpu\":\n    dtype = operand_aval.dtype\n    prefix = \"he\" if dtypes.issubdtype(dtype, np.complexfloating) else \"sy\"\n    target_name = lapack.prepare_lapack_call(f\"{prefix}evd_ffi\",\n                                             operand_aval.dtype)\n    kwargs = {\n      \"mode\": np.uint8(ord(\"V\")),\n      \"uplo\": np.uint8(ord(\"L\" if lower else \"U\")),\n    }\n  else:\n    target_name = f\"{target_name_prefix}solver_syevd_ffi\"\n    # Use Jacobi (algorithm=2) if requested, otherwise use QR (algorithm=1)\n    if algorithm is None:\n      algo_int = 0\n    else:","sourceCodeStart":1275,"sourceCodeEnd":1311,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/linalg.py#L1275-L1311","documentation":"jax/_src/lax/linalg.py:1293 in _eigh_cpu_gpu_lowering. EighImplementation.QDWH (a QR-based Dynamically Weighted Halley iteration used on TPU) has no CPU/GPU implementation. Explicitly selecting algorithm=EighImplementation.QDWH while executing on CPU or GPU raises NotImplementedError.","triggerScenarios":"Calling jax.lax.linalg.eigh(a, algorithm=EighImplementation.QDWH) with backend cpu or gpu. Auto selection (algorithm=None) never picks QDWH off-TPU, so this only occurs with explicit configuration.","commonSituations":"Config copied from a TPU training pipeline to a local GPU/CPU dev machine; experimenting with algorithms for numerical accuracy; library code that hardcodes QDWH for determinism across TPU replicas.","solutions":["Drop the explicit algorithm and use the default per-backend implementation","Select a supported algorithm for CPU/GPU (e.g. EighImplementation.LAPACK-ish default, or JACOBI on GPU)","Run that section on TPU via jax.device_put with tpu backend if QDWH semantics are required"],"exampleFix":"// before\nevals, evecs = jax.lax.linalg.eigh(a, algorithm=lax.linalg.EighImplementation.QDWH)\n// after (on CPU/GPU)\nevals, evecs = jax.lax.linalg.eigh(a)","handlingStrategy":"validation","validationCode":"if jax.default_backend() != 'tpu':\n    algorithm = None  # never force QDWH off-TPU","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Don't hardcode QDWH; auto-select per backend"],"tags":["jax","eigh","qdwh","tpu","backend","not-implemented"],"backgroundTag":"unsupported-operation-on-backend","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}