jax-ml/jax · error · ValueError

Use 'cuda', 'rocm', or 'oneapi' for lax.platform_dependent.

Error message

Use 'cuda', 'rocm', or 'oneapi' for lax.platform_dependent.

What it means

'gpu' is not a valid platform key for lax.platform_dependent; JAX requires the concrete backend name ('cuda', 'rocm', or 'oneapi') because platform-dependent lowering dispatches per concrete compiler.

Source

Thrown at jax/_src/lax/control_flow/conditionals.py:1235

    *args: JAX arrays passed to each of the branches. May be PyTrees.
    **per_platform: branches to use for different platforms. The branches are
      JAX callables invoked with ``*args``. The keywords are platform names,
      e.g., 'cpu', 'tpu', 'cuda', 'rocm'.
    default: optional default branch to use for a platform not mentioned in
      ``per_platform``. If there is no ``default`` there will be an error when
      the code is lowered for a platform not mentioned in ``per_platform``.

  Returns:
    The value ``per_platform[execution_platform](*args)``.
  """
  # Join identical branches
  branches_platforms_list: list[tuple[list[str], Callable]] = []
  for pname, pbranch in per_platform.items():
    if not callable(pbranch):
      raise TypeError(f"lax.platform_dependent: the '{pname}' branch must "
                      "be a callable.")
    if pname == "gpu":
      raise ValueError(
          "Use 'cuda', 'rocm', or 'oneapi' for lax.platform_dependent.")
    for ps, b in branches_platforms_list:
      if b == pbranch:
        ps.append(pname)
        break
    else:
      branches_platforms_list.append(([pname], pbranch))

  platforms_lists, branches = util.unzip2(branches_platforms_list)
  branches_platforms: BranchesPlatforms = tuple(tuple(ps) for ps in platforms_lists)
  if default is not None:
    if not callable(default):
      raise TypeError("lax.platform_dependent: the 'default' branch must "
                      "be a callable.")
    branches = branches + (default,)
    branches_platforms = branches_platforms + (None,)
  platform_index = platform_index_p.bind(platforms=branches_platforms)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Replace 'gpu' with 'cuda' (or 'rocm'/'oneapi' depending on hardware)
If you want one branch for all GPU backends, list them all as separate keys mapping to the same callable (identical branches are merged automatically)
Provide a 'default' branch instead of enumerating backends

Example fix

// before
lax.platform_dependent({'gpu': f_gpu, 'cpu': f_cpu}, x)
// after
lax.platform_dependent({'cuda': f_gpu, 'rocm': f_gpu, 'cpu': f_cpu}, x)

Defensive patterns

Strategy: validation

Validate before calling

VALID = {'cpu', 'cuda', 'rocm', 'tpu', 'oneapi'}
bad = set(per_platform) - VALID
assert not bad, f'use concrete backends, not: {bad}'

Type guard

def valid_platform_keys(per_platform: dict) -> bool:
    return all(k != 'gpu' for k in per_platform)

Try / catch

try: lax.platform_dependent(per_platform, x)\nexcept ValueError as e:\n    if 'cuda' in str(e): per_platform = {**{k: v for k, v in per_platform.items() if k != 'gpu'}, 'cuda': per_platform['gpu'], 'rocm': per_platform['gpu']}\n    else: raise

Prevention

Use concrete backend names: cuda/rocm/oneapi/tpu/cpu
Duplicate the same callable across GPU backend keys; identical branches get merged
Provide a default branch rather than guessing hardware

When it happens

Trigger: Calling lax.platform_dependent({'gpu': gpu_fn, 'cpu': cpu_fn}, ...).

Common situations: Writing platform-dispatch code using the colloquial 'gpu' name; copying code from frameworks that accept 'gpu' as a platform alias.

Related errors

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