jax-ml/jax · error · NotImplementedError

The SVD algorithm parameter is not implemented on TPU.

Error message

The SVD algorithm parameter is not implemented on TPU.

What it means

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.

Source

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

  body_f = lambda args: (
      jnp.array(True),
      jnp.full_like(u_out, np.nan),
      jnp.full_like(s_out, np.nan),
      jnp.full_like(v_out, np.nan),
  )
  _, u_out, s_out, v_out = lax.while_loop(
      cond_f, body_f, (is_finite, u_out, s_out, v_out)
  )

  if is_flip:
    return (v_out, s_out, u_out.T.conj())

  return (u_out, s_out, v_out.T.conj())


def _svd_tpu(a, *, full_matrices, compute_uv, subset_by_index, algorithm=None):
  if algorithm is not None and algorithm != lax_linalg.SvdAlgorithm.DEFAULT:
    raise NotImplementedError(
        "The SVD algorithm parameter is not implemented on TPU.")

  batch_dims = a.shape[:-2]
  fn = functools.partial(
      svd,
      full_matrices=full_matrices,
      compute_uv=compute_uv,
      subset_by_index=subset_by_index,
  )
  for _ in range(len(batch_dims)):
    fn = api.vmap(fn)

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Omit the algorithm argument (DEFAULT works on TPU) or pass SvdAlgorithm.DEFAULT.
2. Conditionally select the algorithm per backend: only set non-DEFAULT on CPU/GPU.
3. If a specific algorithm is required, run that op with jax.default_device(jax.devices('cpu')[0]) or on GPU.

Example fix

# before
u, s, vt = jnp.linalg.svd(a, algorithm=SvdAlgorithm.QR)
# after
u, s, vt = jnp.linalg.svd(a)  # DEFAULT (POLAR) on TPU

import jax
if jax.default_backend() == 'tpu':
    algorithm = None  # DEFAULT only on TPU

try:
    u, s, vt = jnp.linalg.svd(a, algorithm=algorithm)
except NotImplementedError:
    u, s, vt = jnp.linalg.svd(a)  # DEFAULT algorithm fallback

When it happens

Trigger: 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.

Common situations: Tuning algorithm for CPU/GPU performance or stability then running the same code on TPU; cross-backend pipelines promoted to TPU.

Related errors

• Only the POLAR (which is also DEFAULT on TPU) SVD algorithm
• Partitioned callback not implemented on {platform} backend.
• QDWH implementation is only supported on TPU
• Failed to find assignment for logical_axis_index {logical_ax
• masked load_p

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