{"record":{"id":"2d7f0f430d1d9ce1","repo":"jax-ml/jax","slug":"use-cuda-rocm-or-oneapi-for-lax-platform-d","errorCode":null,"errorMessage":"Use 'cuda', 'rocm', or 'oneapi' for lax.platform_dependent.","messagePattern":"Use 'cuda', 'rocm', or 'oneapi' for lax\\.platform_dependent\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/conditionals.py","lineNumber":1235,"sourceCode":"    *args: JAX arrays passed to each of the branches. May be PyTrees.\n    **per_platform: branches to use for different platforms. The branches are\n      JAX callables invoked with ``*args``. The keywords are platform names,\n      e.g., 'cpu', 'tpu', 'cuda', 'rocm'.\n    default: optional default branch to use for a platform not mentioned in\n      ``per_platform``. If there is no ``default`` there will be an error when\n      the code is lowered for a platform not mentioned in ``per_platform``.\n\n  Returns:\n    The value ``per_platform[execution_platform](*args)``.\n  \"\"\"\n  # Join identical branches\n  branches_platforms_list: list[tuple[list[str], Callable]] = []\n  for pname, pbranch in per_platform.items():\n    if not callable(pbranch):\n      raise TypeError(f\"lax.platform_dependent: the '{pname}' branch must \"\n                      \"be a callable.\")\n    if pname == \"gpu\":\n      raise ValueError(\n          \"Use 'cuda', 'rocm', or 'oneapi' for lax.platform_dependent.\")\n    for ps, b in branches_platforms_list:\n      if b == pbranch:\n        ps.append(pname)\n        break\n    else:\n      branches_platforms_list.append(([pname], pbranch))\n\n  platforms_lists, branches = util.unzip2(branches_platforms_list)\n  branches_platforms: BranchesPlatforms = tuple(tuple(ps) for ps in platforms_lists)\n  if default is not None:\n    if not callable(default):\n      raise TypeError(\"lax.platform_dependent: the 'default' branch must \"\n                      \"be a callable.\")\n    branches = branches + (default,)\n    branches_platforms = branches_platforms + (None,)\n  platform_index = platform_index_p.bind(platforms=branches_platforms)\n","sourceCodeStart":1217,"sourceCodeEnd":1253,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/conditionals.py#L1217-L1253","documentation":"'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.","triggerScenarios":"Calling lax.platform_dependent({'gpu': gpu_fn, 'cpu': cpu_fn}, ...).","commonSituations":"Writing platform-dispatch code using the colloquial 'gpu' name; copying code from frameworks that accept 'gpu' as a platform alias.","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"],"exampleFix":"// before\nlax.platform_dependent({'gpu': f_gpu, 'cpu': f_cpu}, x)\n// after\nlax.platform_dependent({'cuda': f_gpu, 'rocm': f_gpu, 'cpu': f_cpu}, x)","handlingStrategy":"validation","validationCode":"VALID = {'cpu', 'cuda', 'rocm', 'tpu', 'oneapi'}\nbad = set(per_platform) - VALID\nassert not bad, f'use concrete backends, not: {bad}'","typeGuard":"def valid_platform_keys(per_platform: dict) -> bool:\n    return all(k != 'gpu' for k in per_platform)","tryCatchPattern":"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","preventionTips":["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"],"tags":["jax","lax","platform-dependent","gpu","invalid-argument"],"backgroundTag":"invalid-platform-name","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}