{"record":{"id":"61e51cfeb40a381d","repo":"jax-ml/jax","slug":"only-the-polar-which-is-also-default-on-tpu-svd","errorCode":null,"errorMessage":"Only the POLAR (which is also DEFAULT on TPU) SVD algorithm is supported on TPU.","messagePattern":"Only the POLAR \\(which is also DEFAULT on TPU\\) SVD algorithm is supported on TPU\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/tpu/linalg/svd.py","lineNumber":291,"sourceCode":"  if compute_uv:\n    u, s, vh = fn(a)\n    return [s, u, vh]\n  else:\n    s = fn(a)\n    return [s]\n\n\ndef _svd_tpu_lowering_rule(\n    ctx, operand, *, full_matrices, compute_uv, subset_by_index, algorithm=None\n):\n  operand_aval, = ctx.avals_in\n  m, n = operand_aval.shape[-2:]\n\n  if algorithm is not None and algorithm not in [\n      lax_linalg.SvdAlgorithm.DEFAULT,\n      lax_linalg.SvdAlgorithm.POLAR,\n  ]:\n    raise NotImplementedError(\n        'Only the POLAR (which is also DEFAULT on TPU) SVD algorithm is'\n        ' supported on TPU.'\n    )\n\n  if m == 0 or n == 0:\n    return mlir.lower_fun(lax_linalg._empty_svd, multiple_results=True)(\n        ctx,\n        operand,\n        full_matrices=full_matrices,\n        compute_uv=compute_uv,\n    )\n\n  return mlir.lower_fun(_svd_tpu, multiple_results=True)(\n      ctx,\n      operand,\n      full_matrices=full_matrices,\n      compute_uv=compute_uv,\n      subset_by_index=subset_by_index,","sourceCodeStart":273,"sourceCodeEnd":309,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/tpu/linalg/svd.py#L273-L309","documentation":"Raised in the MLIR lowering for TPU SVD: only SvdAlgorithm.DEFAULT and SvdAlgorithm.POLAR are supported on TPU (DEFAULT maps to POLAR there). Any other algorithm (QR, JACOBI, etc.) fails at compile/lowering time.","triggerScenarios":"jnp.linalg.svd(..., algorithm=SvdAlgorithm.QR) compiled for TPU, or a jitted function containing a non-POLAR svd lowered to the TPU backend.","commonSituations":"Performance-tuned codeported from GPU to TPU; library code that hard-pins an algorithm; error surfaces during jit/compile rather than at Python call time, making it harder to trace.","solutions":["Remove the explicit algorithm or use SvdAlgorithm.POLAR/DEFAULT on TPU.","Gate algorithm selection on jax.default_backend().","Fall back to another device for that op if the algorithm is essential."],"exampleFix":"# before\nsvd_fn = jax.jit(lambda a: jnp.linalg.svd(a, algorithm=SvdAlgorithm.QR))\n# after\nalg = None if jax.default_backend() == 'tpu' else SvdAlgorithm.QR\nsvd_fn = jax.jit(lambda a: jnp.linalg.svd(a, algorithm=alg))","handlingStrategy":"fallback","validationCode":"SUPPORTED_ON_TPU = {None, SvdAlgorithm.DEFAULT, SvdAlgorithm.POLAR}\nalgorithm = algorithm if (jax.default_backend() != 'tpu' or algorithm in SUPPORTED_ON_TPU) else None","typeGuard":null,"tryCatchPattern":"try:\n    compiled = jax.jit(fn).lower(a)\nexcept NotImplementedError as e:\n    if 'POLAR' in str(e):  # rebuild without algorithm\n        compiled = jax.jit(fn_without_alg).lower(a)\n    else: raise","preventionTips":["Test jit compilation per target backend in CI.","Document TPU-supported algorithm values in project configs."],"tags":["jax","tpu","svd","mlir-lowering","not-implemented"],"backgroundTag":"unsupported-platform-feature","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}