jax-ml/jax · error · TypeError

lax.platform_dependent: the '{pname}' branch must be a calla

Error message

lax.platform_dependent: the '{pname}' branch must be a callable.

What it means

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.

Source

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

  known. This means that the compiler actually never sees a conditional.

  Args:
    *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,)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Wrap each entry in a lambda or function accepting the call args
2. Verify each per_platform value with callable() before calling
3. Use the 'default' parameter for fallback instead of a sentinel value

Example fix

// before
out = lax.platform_dependent({'cpu': x_cpu, 'tpu': x_tpu}, args)
// after
out = lax.platform_dependent({'cpu': lambda a: a * 2, 'tpu': lambda a: a.tpu_op()}, *args)

assert all(callable(b) for b in per_platform.values()), 'branches must be callables'

def valid_platform_branches(per_platform: dict) -> bool:
    return all(isinstance(k, str) and callable(v) for k, v in per_platform.items())

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

When it happens

Trigger: Calling lax.platform_dependent({'cpu': some_array, ...}) or passing a class instance / result value instead of a function per platform key.

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

Related errors

• lax.platform_dependent: the 'default' branch must be a calla
• Use 'cuda', 'rocm', or 'oneapi' for lax.platform_dependent.
• lax.scan: f argument should be a callable.
• lax.while_loop: body_fun and cond_fun arguments should be ca
• iteration over a 0-d array

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