{"record":{"id":"ea2cf861f492e543","repo":"jax-ml/jax","slug":"lax-platform-dependent-the-pname-branch-must","errorCode":null,"errorMessage":"lax.platform_dependent: the '{pname}' branch must be a callable.","messagePattern":"lax\\.platform_dependent: the '(.+?)' branch must be a callable\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/conditionals.py","lineNumber":1232,"sourceCode":"  known. This means that the compiler actually never sees a conditional.\n\n  Args:\n    *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,)","sourceCodeStart":1214,"sourceCodeEnd":1250,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/conditionals.py#L1214-L1250","documentation":"lax.platform_dependent requires each per-platform entry to be a callable (a function taking the same args), not a precomputed value or other object. Passing anything non-callable raises this TypeError at trace time.","triggerScenarios":"Calling lax.platform_dependent({'cpu': some_array, ...}) or passing a class instance / result value instead of a function per platform key.","commonSituations":"Developers assuming branches are values like in a dict lookup; migrating code from Python if/else on jax.default_backend() and passing computed results directly.","solutions":["Wrap each entry in a lambda or function accepting the call args","Verify each per_platform value with callable() before calling","Use the 'default' parameter for fallback instead of a sentinel value"],"exampleFix":"// before\nout = lax.platform_dependent({'cpu': x_cpu, 'tpu': x_tpu}, args)\n// after\nout = lax.platform_dependent({'cpu': lambda a: a * 2, 'tpu': lambda a: a.tpu_op()}, *args)","handlingStrategy":"type-guard","validationCode":"assert all(callable(b) for b in per_platform.values()), 'branches must be callables'","typeGuard":"def valid_platform_branches(per_platform: dict) -> bool:\n    return all(isinstance(k, str) and callable(v) for k, v in per_platform.items())","tryCatchPattern":"try: lax.platform_dependent(per_platform, *args)\\nexcept TypeError as e:\\n    if 'must be a callable' in str(e): wrap values in lambdas and retry\\n    else: raise","preventionTips":["Always write branches as 'name': lambda args: ...","Lint for non-callable dict values passed to platform_dependent","Remember branches receive the call arguments lazily"],"tags":["jax","lax","platform-dependent","typeerror","callable"],"backgroundTag":"callback-must-be-callable","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}