code: def _jit_update_fun(predict_fun, loss_fun, optimizer, lr_fun, num_devices):
    
    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
docs: Get jit-ed update function for loss, optimizer, learning rate function.