def mapped_update(weights_and_slots, i, opt_params, batch, state, rng):
    """This is a multi-device version of the update function above."""
    # We assume all tensors have the first dimension = n_devices.
    weights, slots = weights_and_slots
    rng, subrng = jax_random.split(rng)
    grad_fn = fastmath.grad(model_and_loss_call, has_aux=True)
    grads, state = grad_fn(weights, batch, state, rng)
    # We do a psum(1.0) here instead of `n_devices` since `n_devices` is just
    # the number of devices on this host machine, however psum goes over all
    # devices of all hosts (ex: a TPU pod) and we need to be averaging over all
    # of them.
    #
    # Collect all gradients.
    grads = fastmath.psum(grads, 'batch')
    n_devices_total = fastmath.psum(np.array(1.0), 'batch')
    # Average across hosts.
    grads = jax.tree_util.tree_map(lambda g: g / n_devices_total, grads)

    new_weights, new_slots, stats = optimizer.tree_update(
        i, grads, weights, slots, opt_params)
    return (new_weights, new_slots), stats, state, subrng
This is a multi-device version of the update function above.
trax/supervised/trainer_lib.py