code: def _get_jit_or_pmap_broadcast():
      """Initialize a cache of pmapped functions closed over non-`jnp.ndarray` args.
  
      Returns:
        A `jit_or_pmap_broadcast` function allowing to jit or pmap a function as a
        closure over all non-`jnp.ndarray` args, all `kwargs`, while broadcasting
        all `jnp.ndarray`s in `args` except the first one.
      """
      cache = {}
  
      def jit_or_pmap_broadcast(
          f: Callable, device_count: int = -1
      ) -> Callable:
          """Pmap `f` over the first argument by closing over or broadcasting others.
  
          Args:
            f:
              function to pmap. First argument must be an `jnp.ndarray` or a Kernel.
              In either case, ndarrays should have a leading axis having the size of
              `device_count`.
  
            device_count:
              number of XLA devices. `-1` means all available devices. `0` means to
              just `jit` the function.
  
          Returns:
            A function of the same signature as `f` pmapped over the `jnp.ndarray`s in
            the first argument. Other arguments are either closed over
            (non-`jnp.ndarray`s in `args` and all `kwargs`) or broadcasted to
            `(device_count,) + old_shape` (for `jnp.ndarray`s). If
            `device_count == 0`, `f` is closed over and jitted over all non-array
            arguments and all `kwargs`.
  
          Raises:
            An error if `kwargs` have a `jnp.ndarray`.
            TODO(romann): treat `jnp.ndarray`s in `kwargs` when JAX allows it. See
            https://github.com/google/jax/issues/912
          """
          key = (f, device_count)
  
          if device_count == -1:
              device_count = jax.local_device_count()
  
          # TODO(romann): adapt this when JAX allows `axis_in` for `pmap`.
          def broadcast(arg: jnp.ndarray) -> jnp.ndarray:
              if device_count == 0:
                  return arg
              return jnp.broadcast_to(arg, (device_count,) + arg.shape)
  
          @utils.wraps(f)
          def f_pmapped(
              x_or_kernel: Union[jnp.ndarray, Kernel], *args, **kwargs
          ):
              args_np, args_np_idxs = [], []
              args_other = {}
  
              # TODO(romann): treat `jnp.ndarray`s in `kwargs` when JAX allows it.
              # https://github.com/google/jax/issues/912
              # Filter out `jnp.ndarray`s from other arguments.
              for i, arg in enumerate(args):
                  if _is_np_ndarray(arg):
                      args_np.append(arg)
                      args_np_idxs.append(i)
                  else:
                      args_other[i] = arg
              kwargs_np = {}
              kwargs_other = {}
              for k, v in kwargs.items():
                  if _is_np_ndarray(v):
                      assert isinstance(v, tuple), len(v) == 2
                      kwargs_np[k] = (v[0], broadcast(v[1]))
                  else:
                      kwargs_other[k] = v
  
              # Check cache before jitting.
              _key = key + (
                  tuple(args_other.items())
                  + tuple(kwargs_other.items())
              )
              if _key in cache:
                  _f = cache[_key]
              else:
                  # Define a `jnp.ndarray`-only function as a closure over other args.
                  def _f(_x_or_kernel, *_args_np, **_kwargs_np):
                      # Merge args.
                      _args_np = {
                          i: _arg_np
                          for i, _arg_np in zip(args_np_idxs, _args_np)
                      }
                      _args = {**_args_np, **args_other}
                      _args = tuple(v for k, v in sorted(_args.items()))
                      _kwargs = {**_kwargs_np, **kwargs_other}
                      return f(_x_or_kernel, *_args, **_kwargs)
  
                  _f = jit(_f) if device_count == 0 else pmap(_f)
                  cache[_key] = _f
  
              # Broadcast `jnp.ndarray` arguments and apply the new function to them.
              args_np = tree_map(broadcast, args_np)
              return _f(x_or_kernel, *args_np, **kwargs_np)
  
          return f_pmapped
  
      return jit_or_pmap_broadcast
  
imports: from functools import partial
  from typing import Any, Callable, Iterable, Optional, TypeVar, Union
  import warnings
  import jax
  from jax import device_put
  from jax import devices
  from jax import jit
  from jax import pmap
  from jax import random
  import jax.numpy
  from jax.tree_util import tree_all
  from jax.tree_util import tree_flatten
  from jax.tree_util import tree_map
  from jax.tree_util import tree_unflatten
  import numpy
  from .utils import utils
  from .utils.kernel import Kernel
  from .utils.typing import KernelFn
  from .utils.typing import NTTree
label_desc: Use image node_id 12 for calling a global function with example usage: _get_jit_or_pmap_broadcast() and returns: jit_or_pmap_broadcast, f_pmapped, jnp, _f, arg, f