jax-ml/jax · error · NotImplementedError

Only the POLAR (which is also DEFAULT on TPU) SVD algorithm

Error message

Only the POLAR (which is also DEFAULT on TPU) SVD algorithm is supported on TPU.

What it means

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.

Source

Thrown at jax/_src/tpu/linalg/svd.py:291

  if compute_uv:
    u, s, vh = fn(a)
    return [s, u, vh]
  else:
    s = fn(a)
    return [s]


def _svd_tpu_lowering_rule(
    ctx, operand, *, full_matrices, compute_uv, subset_by_index, algorithm=None
):
  operand_aval, = ctx.avals_in
  m, n = operand_aval.shape[-2:]

  if algorithm is not None and algorithm not in [
      lax_linalg.SvdAlgorithm.DEFAULT,
      lax_linalg.SvdAlgorithm.POLAR,
  ]:
    raise NotImplementedError(
        'Only the POLAR (which is also DEFAULT on TPU) SVD algorithm is'
        ' supported on TPU.'
    )

  if m == 0 or n == 0:
    return mlir.lower_fun(lax_linalg._empty_svd, multiple_results=True)(
        ctx,
        operand,
        full_matrices=full_matrices,
        compute_uv=compute_uv,
    )

  return mlir.lower_fun(_svd_tpu, multiple_results=True)(
      ctx,
      operand,
      full_matrices=full_matrices,
      compute_uv=compute_uv,
      subset_by_index=subset_by_index,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the explicit algorithm or use SvdAlgorithm.POLAR/DEFAULT on TPU.
2. Gate algorithm selection on jax.default_backend().
3. Fall back to another device for that op if the algorithm is essential.

Example fix

# before
svd_fn = jax.jit(lambda a: jnp.linalg.svd(a, algorithm=SvdAlgorithm.QR))
# after
alg = None if jax.default_backend() == 'tpu' else SvdAlgorithm.QR
svd_fn = jax.jit(lambda a: jnp.linalg.svd(a, algorithm=alg))

SUPPORTED_ON_TPU = {None, SvdAlgorithm.DEFAULT, SvdAlgorithm.POLAR}
algorithm = algorithm if (jax.default_backend() != 'tpu' or algorithm in SUPPORTED_ON_TPU) else None

try:
    compiled = jax.jit(fn).lower(a)
except NotImplementedError as e:
    if 'POLAR' in str(e):  # rebuild without algorithm
        compiled = jax.jit(fn_without_alg).lower(a)
    else: raise

When it happens

Trigger: jnp.linalg.svd(..., algorithm=SvdAlgorithm.QR) compiled for TPU, or a jitted function containing a non-POLAR svd lowered to the TPU backend.

Common situations: 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.

Related errors

• The SVD algorithm parameter is not implemented on TPU.
• QDWH implementation is only supported on TPU
• Failed to find assignment for logical_axis_index {logical_ax
• masked load_p
• run_scoped_p with collective axes is not supported

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/61e51cfeb40a381d. Report an issue: GitHub.