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)
Averages gradients over all the devices across different hosts.
trax/optimizers/trainer.py