def _jit_update_fn(predict_fn, loss_fn, optimizer, n_devices, jit=True):
  """Returns a (JIT-compiled) function that computes updates for one step."""
  model_and_loss = tl.Serial(predict_fn, loss_fn)
  # Gradients are always wrt. the first argument, so putting weights first.
  def model_and_loss_call(weights, batch, state, rng):
    res = model_and_loss(batch, weights=weights, state=state, rng=rng)
    return res, model_and_loss.state
  if n_devices == 1:  # TODO(lukaszkaiser): remove branch when not needed.
    def single_update(weights_and_slots, i, opt_params, batch, state, rng):
      weights, slots = weights_and_slots
      rng, subrng = jax_random.split(rng[0])
      grad_fn = fastmath.grad(model_and_loss_call, has_aux=True)
      grads, state = grad_fn(weights, batch, state, rng)
      new_weights, new_slots, stats = optimizer.tree_update(
          i, grads, weights, slots, opt_params)
      return (new_weights, new_slots), stats, state, [subrng]
    if jit:
      # TODO(lukaszkaiser): donate_argnums=(0,) when XLA supports it on GPU
      return fastmath.jit(single_update)
    else:
      return single_update

  # Else, for n_devices > 1:
  @functools.partial(fastmath.pmap, axis_name='batch')  # donate_argnums=(0,))
  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

  def update(weights_and_slots, i, opt_params, batch, state, rng):
    return mapped_update(weights_and_slots, np.repeat(i, n_devices),
                         opt_params, batch, state, rng)

  return update
Returns a (JIT-compiled) function that computes updates for one step.
trax/supervised/trainer_lib.py