code: def broadcast_one_to_all(
      in_tree: Any, is_source: bool | None = None
  ) -> Any:
      """Broadcast data from a source host (host 0 by default) to all other hosts.
  
      Args:
        in_tree: pytree of arrays - each array *must* have the same shape across the
          hosts.
        is_source: optional bool denoting whether the caller is the source. Only
          'source host' will contribute the data for the broadcast. If None, then
          host 0 is used.
  
      Returns:
        A pytree matching in_tree where the leaves now all contain the data from the
        first host.
      """
      if is_source is None:
          is_source = jax.process_index() == 0
  
      devices: np.ndarray = np.array(jax.devices()).reshape(
          jax.process_count(), jax.local_device_count()
      )
      global_mesh = jax.sharding.Mesh(
          devices, ("processes", "local_devices")
      )
      pspec = P("processes")
  
      def pre_jit(x):
          if is_source:
              inp = x
          else:
              inp = np.zeros_like(x)
          inp = np.expand_dims(inp, axis=0)
          return host_local_array_to_global_array(
              inp, global_mesh, pspec
          )
  
      def post_jit(x):
          return np.asarray(x.addressable_data(0))
  
      in_tree = jax.tree_map(pre_jit, in_tree)
      out_tree = jax.jit(
          _psum,
          out_shardings=jax.sharding.NamedSharding(global_mesh, P()),
      )(in_tree)
      return jax.tree_map(post_jit, out_tree)
  
imports: from __future__ import annotations
  from functools import partial, lru_cache
  from typing import Optional
  import zlib
  from typing import Any
  import jax
  import jax.numpy
  from jax.tree_util import tree_flatten, tree_map, tree_unflatten
  from jax._src import core
  from jax._src.interpreters import ad
  from jax._src.interpreters import batching
  from jax._src.interpreters import mlir
  from jax._src import array
  from jax._src import sharding_impls
  from jax._src.interpreters import pxla
  from jax.interpreters import xla
  from jax._src import pjit
  from jax.experimental.pjit import pjit
  from jax.sharding import PartitionSpec
  from jax._src import distributed
  from jax._src.util import safe_zip
  import numpy
label_desc: Use image node_id 2 for calling a global function with example usage: broadcast_one_to_all(in_tree, is_source) and returns: jax, host_local_array_to_global_array, np