code: 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
  
imports: import collections
  import functools
  import itertools
  import os
  import sys
  import time
  from absl import logging
  import gin
  import jax
  import tensorflow.compat.v2
  from trax import fastmath
  from trax import jaxboard
  from trax import layers
  from trax import optimizers
  from trax.data import inputs
  from trax.fastmath import numpy
  from trax.fastmath import random
  from trax.layers import base
  from trax.shapes import ShapeDtype
  from trax.supervised import history
  from trax.supervised import lr_schedules
  from trax.supervised import training
label_desc: Use image node_id 4 for calling a global function with example usage: _jit_update_fn(predict_fn, loss_fn, optimizer, n_devices, jit) and returns: update, res,  model_and_loss, new_weights,  new_slots,  stats,  state,  subrng, mapped_update, new_weights,  new_slots,  stats,  state,  list, fastmath, single_update