{"record":{"id":"bfbb39761785e713","repo":"jax-ml/jax","slug":"the-svd-algorithm-parameter-is-not-implemented-on","errorCode":null,"errorMessage":"The SVD algorithm parameter is not implemented on TPU.","messagePattern":"The SVD algorithm parameter is not implemented on TPU\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/tpu/linalg/svd.py","lineNumber":260,"sourceCode":"  body_f = lambda args: (\n      jnp.array(True),\n      jnp.full_like(u_out, np.nan),\n      jnp.full_like(s_out, np.nan),\n      jnp.full_like(v_out, np.nan),\n  )\n  _, u_out, s_out, v_out = lax.while_loop(\n      cond_f, body_f, (is_finite, u_out, s_out, v_out)\n  )\n\n  if is_flip:\n    return (v_out, s_out, u_out.T.conj())\n\n  return (u_out, s_out, v_out.T.conj())\n\n\ndef _svd_tpu(a, *, full_matrices, compute_uv, subset_by_index, algorithm=None):\n  if algorithm is not None and algorithm != lax_linalg.SvdAlgorithm.DEFAULT:\n    raise NotImplementedError(\n        \"The SVD algorithm parameter is not implemented on TPU.\")\n\n  batch_dims = a.shape[:-2]\n  fn = functools.partial(\n      svd,\n      full_matrices=full_matrices,\n      compute_uv=compute_uv,\n      subset_by_index=subset_by_index,\n  )\n  for _ in range(len(batch_dims)):\n    fn = api.vmap(fn)\n\n  if compute_uv:\n    u, s, vh = fn(a)\n    return [s, u, vh]\n  else:\n    s = fn(a)\n    return [s]","sourceCodeStart":242,"sourceCodeEnd":278,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/tpu/linalg/svd.py#L242-L278","documentation":"The TPU implementation of SVD only supports the DEFAULT algorithm; explicitly requesting a different SvdAlgorithm (e.g. QR or JACOBI_ANNIHILATOR, which select CPU/GPU algorithms) raises NotImplementedError in the TPU frontend wrapper.","triggerScenarios":"Calling jax.numpy.linalg.svd(a, algorithm=lax_linalg.SvdAlgorithm.QR) (or any non-DEFAULT value) on a TPU backend, via jax._src.tpu.linalg.svd._svd_tpu.","commonSituations":"Tuning algorithm for CPU/GPU performance or stability then running the same code on TPU; cross-backend pipelines promoted to TPU.","solutions":["Omit the algorithm argument (DEFAULT works on TPU) or pass SvdAlgorithm.DEFAULT.","Conditionally select the algorithm per backend: only set non-DEFAULT on CPU/GPU.","If a specific algorithm is required, run that op with jax.default_device(jax.devices('cpu')[0]) or on GPU."],"exampleFix":"# before\nu, s, vt = jnp.linalg.svd(a, algorithm=SvdAlgorithm.QR)\n# after\nu, s, vt = jnp.linalg.svd(a)  # DEFAULT (POLAR) on TPU","handlingStrategy":"fallback","validationCode":"import jax\nif jax.default_backend() == 'tpu':\n    algorithm = None  # DEFAULT only on TPU","typeGuard":null,"tryCatchPattern":"try:\n    u, s, vt = jnp.linalg.svd(a, algorithm=algorithm)\nexcept NotImplementedError:\n    u, s, vt = jnp.linalg.svd(a)  # DEFAULT algorithm fallback","preventionTips":["Backend-gate any non-DEFAULT algorithm choice.","Keep algorithm as a config knob defaulted to None."],"tags":["jax","tpu","svd","not-implemented","backend-support"],"backgroundTag":"unsupported-platform-feature","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}