{"record":{"id":"6847e48232acd5dc","repo":"jax-ml/jax","slug":"subset-by-index-not-supported-on-cpu-and-gpu","errorCode":null,"errorMessage":"subset_by_index not supported on CPU and GPU","messagePattern":"subset_by_index not supported on CPU and GPU","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/linalg.py","lineNumber":1289,"sourceCode":"    )\n  n = shape[0]\n  d = (n if subset_by_index is None else\n       subset_by_index[1] - subset_by_index[0])\n  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\"","sourceCodeStart":1271,"sourceCodeEnd":1307,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/linalg.py#L1271-L1307","documentation":"jax/_src/lax/linalg.py:1289 in _eigh_cpu_gpu_lowering. The subset_by_index argument of jax.lax.linalg.eigh (computing only eigenpairs [lo, hi)) is only implemented on TPU. On CPU and GPU the lowering only accepts the full range, so any other subset raises NotImplementedError.","triggerScenarios":"Calling jax.lax.linalg.eigh(a, subset_by_index=(lo, hi)) with a range other than (0, n) or None while running on CPU or GPU backend. Note lower bounds other than 0 are also unsupported even on TPU.","commonSituations":"Porting scipy.linalg.eigh(..., subset_by_index=...) partial-spectrum code to JAX; memory-saving attempts to compute only the top-k eigenvalues of large matrices; code that ran on TPU moved to a GPU box.","solutions":["Compute the full spectrum and slice: evals, evecs = eigh(a); take evals[lo:hi] (fine when n is moderate)","For top-k on GPU/CPU use an iterative method (e.g. jax.scipy.sparse.linalg.lobpcg) or scipy.sparse.linalg.eigsh on the host","Run on TPU backend where subset_by_index is supported"],"exampleFix":"// before\nevals, evecs = jax.lax.linalg.eigh(a, subset_by_index=(0, 10))  # on GPU\n// after\nevals_all, evecs_all = jax.lax.linalg.eigh(a)\nevals, evecs = evals_all[..., :10], evecs_all[..., :, :10]","handlingStrategy":"fallback","validationCode":"backend = jax.default_backend()\nuse_subset = backend == 'tpu' and lo == 0\nif not use_subset:\n    plan = 'full eigh + slice'","typeGuard":null,"tryCatchPattern":"try:\n    w, v = jax.lax.linalg.eigh(a, subset_by_index=(lo, hi))\nexcept NotImplementedError:\n    w_all, v_all = jax.lax.linalg.eigh(a)\n    w, v = w_all[..., lo:hi], v_all[..., :, lo:hi]","preventionTips":["Remember subset_by_index is TPU-only","Slice full results on CPU/GPU"],"tags":["jax","eigh","subset-by-index","cpu","gpu","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"}