code: def platform_dependent(
      *args: Any,
      default: Callable[..., _T] | None = None,
      **per_platform: Callable[..., _T],
  ):
      """Stages out platform-specific code.
  
      In JAX the actual platform on which a computation is run is determined
      very late, e.g., based on where the data is located. When using AOT
      lowering or serialization, the computation may be compiled and executed
      on a different machine, or even on a platform that is not available at
      lowering time. This means that it is not safe to write platform-dependent
      code using Python conditionals, e.g., based on the current default
      JAX platform. Instead, one can use ``platform_dependent``:
  
      Usage::
  
          def cpu_code(*args): ...
          def tpu_code(*args): ...
          def other_platforms_code(*args): ...
          res = platform_dependent(*args, cpu=cpu_code, tpu=tpu_code,
                                   default=other_platforms_code)
  
      When the staged out code is executed on a CPU, this is equivalent to
      ``cpu_code(*args)``, on a TPU is equivalent to ``tpu_code(*args)`` and on
      any other platform to ``other_platforms_code(*args)``.
      Unlike a Python conditional, all alternatives are traced
      and staged out to Jaxpr. This is similar to, and is implemented in terms of,
      :func:`~switch`, from which it inherits the behavior
      under transformations.
  
      Unlike a :func:`~switch` the choice of what gets executed is made earlier:
      in most cases during lowering when the lowering platform is known; in the
      rare case of multi-platform lowering and serialization, the StableHLO code
      will contain a conditional on the actual platform. This conditional is
      resolved just in time prior to compilation when the compilation platform is
      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
      platform_branches: list[tuple[list[str], Callable]] = []
      for pname, pbranch in per_platform.items():
          if pname == "gpu":
              raise ValueError(
                  "Use 'cuda' or 'rocm' for lax.platform_dependent."
              )
          for ps, b in platform_branches:
              if b == pbranch:
                  ps.append(pname)
                  break
          else:
              platform_branches.append(([pname], pbranch))
  
      platforms_lists, branches = util.unzip2(platform_branches)
      platform_index = platform_index_p.bind(
          platforms=tuple(tuple(ps) for ps in platforms_lists),
          has_default=(default is not None),
      )
      if default is not None:
          branches = branches + (default,)
      # Use a switch, to get the proper transformation rules for free. Since
      # platform index has no dependence on the input data, it won't be vectorized
      # under vmap.
      return switch(platform_index, branches, *args)
  
imports: from __future__ import annotations
  import collections
  from collections.abc import Sequence
  import functools
  from functools import partial
  import inspect
  import itertools
  import operator
  from typing import Any, Callable, TypeVar
  from jax.tree_util import tree_flatten, tree_unflatten
  from jax._src import ad_util
  from jax._src import config
  from jax._src import core
  from jax._src import dispatch
  from jax._src import dtypes
  from jax._src import effects
  from jax._src import linear_util
  from jax._src import source_info_util
  from jax._src import util
  from jax._src.state.discharge import register_discharge_rule, discharge_state
  from jax._src.state.types import AbstractRef, RefEffect
  from jax._src.core import ConcreteArray, raise_to_shaped, replace_jaxpr_effects
  from jax._src.interpreters import ad
  from jax._src.interpreters import batching
  from jax._src.interpreters import mlir
  from jax._src.interpreters import partial_eval
  from jax._src.interpreters import xla
  from jax._src.lax import lax
  from jax._src.traceback_util import api_boundary
  from jax._src.util import safe_map, split_list, partition_list
  from jax._src.lib.mlir import ir
  from jax._src.lib.mlir.dialects import hlo
  import numpy
  from jax._src.lax.control_flow.common import _abstractify, _avals_short, _check_tree_and_avals, _initial_style_jaxprs_with_common_consts, _make_closed_jaxpr, _prune_zeros, _typecheck_param
label_desc: Use image node_id 25 for calling a global function with example usage: platform_dependent() and returns: switch