def _make_weights_and_state_same_across_hosts(weights_and_state):
  """Makes train and eval model's weights and state the same across hosts."""

  # We assume that weights_and_state have been already replicated, i.e the
  # leading axis is self._n_devices

  # This is the total number of devices across all hosts.
  n_devices_total = fastmath.psum(jnp.array(1.0), 'devices').astype(jnp.int32)

  # We average the weights and state across all devices.
  # We also make sure we don't change the type of the weights and state.
  return fastmath.nested_map(
      lambda x: (fastmath.psum(x, 'devices') / n_devices_total).astype(x.dtype),
      weights_and_state)
Makes train and eval model's weights and state the same across hosts.
trax/supervised/training.py