code: def make_fit_op(
      loss_fn,
      optimizer,
      trainable_variables,
      grad_summary_fn=None,
      tf_function=True,
      xla_compile=True,
  ):
      """One training step.
  
      Args:
        loss_fn: Python `callable` which returns the pair `loss` (`tf.Tensor`) and
          any other second result such that
          `tf.nest.map_structure(tf.convert_to_tensor, other)` will succeed.
        optimizer: `tf_keras.optimizers.Optimizer`-like instance which has members
          `gradient` and `apply_gradients`.
        trainable_variables: `tf.nest.flatten`-able structure of `tf.Variable`
          instances.
        grad_summary_fn: Python `callable` which takes a `trainable_variables`-like
          structure of `tf.Tensor`s representing the gradient of the result of
          `loss_fn` with respect to `trainable_variables`. For example,
          `lambda grads: tf.nest.map_structure(
             lambda x: 0. if x is None else tf.norm(x), grads)`.
          Default value: `None` (i.e., no summarization is made).
        tf_function: `bool` representing whether the resulting function should be
          `tf.function` decoreated.
          Default value: `True`.
        xla_compile: `bool` representing whether XLA compilation should be
          performed. (This argument is ignored if the function is executed eagerly.)
          Default value: `True`.
  
      Returns:
        fit_op: A Python `callable` taking args which are forwarded to `loss_fn` and
          such that when called `trainable_variables` are updated per the logic of
          `optimizer.apply_gradients`.
      """
  
      @tfcompile(tf_function=tf_function, xla_best_effort=xla_compile)
      def fit_op(*args, **kwargs):
          """Performs one gradient descent update to `trainable_variables`."""
          with tf.GradientTape(watch_accessed_variables=False) as tape:
              tf.nest.map_structure(tape.watch, trainable_variables)
              loss, other = loss_fn(*args, **kwargs)
          grads = tf.nest.pack_sequence_as(
              trainable_variables,
              tape.gradient(loss, tf.nest.flatten(trainable_variables)),
          )
          if isinstance(optimizer, collections.abc.Sequence):
              for opt, g, v in zip(
                  optimizer, grads, trainable_variables
              ):
                  _apply_gradients(opt, g, v)
          else:
              _apply_gradients(optimizer, grads, trainable_variables)
          if grad_summary_fn is not None:
              return loss, other, grad_summary_fn(grads)
          return loss, other
  
      # Note: we can't do `return tf.xla.experimental.compile(fit)` since we can't
      # assume the function arguments are coercible to `tf.Tensor`s.
      return fit_op
  
imports: import collections
  import contextlib
  import functools
  import sys
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
label_desc: Use image node_id 5 for calling a global function with example usage: make_fit_op(loss_fn, optimizer, trainable_variables, grad_summary_fn, tf_function, xla_compile) and returns: fit_op, loss,  other, loss,  other,  grad_summary_fn