Get jit-ed update function for loss, optimizer, learning rate function.

def _jit_update_fun(predict_fun, loss_fun, optimizer, lr_fun, num_devices):
  """Get jit-ed update function for loss, optimizer, learning rate function."""
  if num_devices == 1:  # TODO(lukaszkaiser): remove branch when not needed.
    def single_update(i, opt_state, batch, rng):
      rng, subrng = jax_random.split(rng[0])
      _, opt_update = optimizer(lr_fun)
      params = trax_opt.get_params(opt_state)
      return opt_update(i, backend.grad(loss_fun)(
          params, batch, predict_fun, rng), opt_state), [subrng]
    return backend.jit(single_update)

  @functools.partial(backend.pmap, axis_name="batch")
  def mapped_update(i, opt_state, batch, rng):
    """This is a multi-device version of the update function above."""
    # We assume all tensors have the first dimension = num_devices.
    rng, subrng = jax_random.split(rng)
    _, opt_update = optimizer(lr_fun)
    params = trax_opt.get_params(opt_state)
    grads = backend.grad(loss_fun)(params, batch, predict_fun, rng)
    grads = jax.tree_util.tree_map(
        lambda g: lax.psum(g, "batch"), grads)
    return opt_update(i, grads, opt_state), subrng

  def update(i, opt_state, batch, rng):
    return mapped_update(jax.replicate(i), opt_state, batch, rng)

  return update