code: def host_local_array_to_global_array_impl(
      arr: Any, *, global_mesh: jax.sharding.Mesh, pspec: Any
  ):
      if pspec is None:
          raise ValueError(
              "`None` is not a valid input to the pspecs argument. Please use "
              "jax.sharding.PartitionSpec() if you wanted to replicate your input."
          )
      # If the Array is not fully addressable i.e. not host local, return it.
      if (
          isinstance(arr, array.ArrayImpl)
          and not arr.is_fully_addressable
      ):
          return arr
      if isinstance(arr, array.ArrayImpl) and isinstance(
          arr.sharding, jax.sharding.PmapSharding
      ):
          arr = np.array(arr)
  
      local_sharding = jax.sharding.NamedSharding(
          global_mesh.local_mesh, pspec
      )
  
      # If the input is a concrete jax.Array and the input array sharding
      # matches the `local_sharding`, then there's no need to reshard and create
      # copies.
      if isinstance(
          arr, array.ArrayImpl
      ) and arr.sharding.is_equivalent_to(local_sharding, arr.ndim):
          arrays = [x.data for x in arr.addressable_shards]
      else:
          arr = xla.canonicalize_dtype(arr)
          arrays = [
              arr[index]
              for d, index in local_sharding.devices_indices_map(
                  arr.shape
              ).items()
          ]
  
      global_aval = _local_to_global_aval(
          core.ShapedArray(arr.shape, arr.dtype), global_mesh, pspec
      )
  
      return pxla.batched_device_put(
          global_aval,
          jax.sharding.NamedSharding(global_mesh, pspec),
          arrays,
          list(global_mesh.local_mesh.devices.flat),
      )
  
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 12 for calling a global function with example usage: host_local_array_to_global_array_impl(arr) and returns: pxla, arr