code: def _average_multidevice_gradients(gradients, adasum=False):
      """Averages gradients over all the devices across different hosts."""
      n = fastmath.global_device_count() // base.N_WEIGHTS_SHARDS
      if adasum:
          # This implements a version of the Adasum algorithm from the following
          # paper: https://arxiv.org/pdf/2006.02924.pdf
          lg = max([i for i in range(20) if 2**i <= n])
          for lg_i in range(lg):
              shift = 2**lg_i
              perm = []
              for i in range(n):
                  block_i = i % (
                      2 * shift
                  )  # we do blocks of 2*shift size
                  if block_i < shift:
                      perm.append((i, i + shift))
                  else:
                      perm.append((i, i - shift))
              perm_grad = jax.lax.ppermute(
                  gradients, perm=perm, axis_name="batch"
              )
              gradients = fastmath.nested_map_multiarg(
                  _adasum_merge, gradients, perm_grad
              )
      if (
          base.N_WEIGHTS_SHARDS > 1
      ):  # only sum gradients from matching shards
          groups = [
              [base.N_WEIGHTS_SHARDS * i + d for i in range(int(n))]
              for d in range(base.N_WEIGHTS_SHARDS)
          ]
          gradients_psum = fastmath.psum(
              gradients, "batch", axis_index_groups=groups
          )
      else:
          gradients_psum = fastmath.psum(
              gradients, "batch"
          )  # sum all gradients
      n = jnp.array(n, dtype=jnp.float32)
      return fastmath.nested_map(lambda g: g / n, gradients_psum)
  
imports: from concurrent import futures
  import functools
  import os
  import time
  from absl import logging
  import jax
  import numpy
  import psutil
  from trax import fastmath
  from trax import layers
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import combinators
label_desc: Use image node_id 2 for calling a global function with example usage: _average_multidevice_gradients(gradients, adasum) and returns: fastmath