jax-ml/jax · error · TypeError

lax.platform_dependent: the 'default' branch must be a calla

Error message

lax.platform_dependent: the 'default' branch must be a callable.

What it means

The optional 'default' argument of lax.platform_dependent must be a callable just like the named platform branches. Passing a value or None-adjacent sentinel that isn't callable raises this TypeError.

Source

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

  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)

  if core.is_concrete(platform_index):
    return branches[int(platform_index)](*args)
  return _switch_internal(platform_index, branches, args,
                          branches_platforms=branches_platforms)


# A primitive to compute the index of a platform into a list of platforms.
# Args:
#   platforms: BranchesPlatforms. If the current lowering
#     platform is in one of the inner tuples returns the index of that inner
#     tuple in the outer tuple.
platform_index_p = core.Primitive("platform_index")
platform_index_p.multiple_results = False

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass default as a zero/one-arg function returning the fallback computation
2. Precompute nothing: let the function receive args and compute lazily
3. Check callable(default) before the call in generic wrappers

Example fix

// before
lax.platform_dependent({'cuda': f}, x, default=x * 2)
// after
lax.platform_dependent({'cuda': f}, x, default=lambda a: a * 2)

if default is not None and not callable(default):
    val = default
    default = (lambda *a: val)  # wrap value in a callable

def default_is_callable(default) -> bool:
    return default is None or callable(default)

try: lax.platform_dependent(p, x, default=d)
except TypeError as e:
    if 'default' in str(e): d = (lambda *a: d); retry
    else: raise

When it happens

Trigger: Calling lax.platform_dependent(per_platform, *args, default=some_array) or default=some_object.

Common situations: Assuming default is a fallback value rather than a fallback function; reusing a computed result as the default.

Related errors

• lax.platform_dependent: the '{pname}' 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/9d6f0c8352c9cfc8. Report an issue: GitHub.