code: def per_host_sum_pmap(in_tree):
      """Execute psum on in_tree"s leaves over one device per host."""
      host2devices = collections.defaultdict(list)
      for d in jax.devices():
          host2devices[d.process_index].append(d)
      devices = [host2devices[k][0] for k in host2devices]
      host_psum = jax.pmap(
          lambda x: jax.lax.psum(x, "i"), "i", devices=devices
      )
  
      def pre_pmap(xs):
          return jax.tree_util.tree_map(
              lambda x: jnp.broadcast_to(x, (1,) + x.shape), xs
          )
  
      def post_pmap(xs):
          return jax.tree_util.tree_map(lambda x: x[0], xs)
  
      return post_pmap(host_psum(pre_pmap(in_tree)))
  
imports: import collections
  import os
  from absl import logging
  from clu import metric_writers
  from clu import periodic_actions
  from flax import linen
  from flax.training import checkpoints
  from flax.training import common_utils
  import jax
  from jax import random
  import jax.numpy
  from jax.sharding import PartitionSpec, Mesh, NamedSharding
  import ml_collections
  import numpy
  import optax
  import tensorflow
  import input_pipeline
  import models
  import temperature_sampler
  import utils
label_desc: Use image node_id 10 for calling a global function with example usage: per_host_sum_pmap(in_tree) and returns: post_pmap, jax, jax