code: def _make_training_loop_body(
      optimizer_step_fn,
      convergence_criterion,
      batch_convergence_reduce_fn,
      seed_is_none=False,
  ):
      """Constructs the training loop body."""
  
      def unconverged_training_loop_body(previous, step):
          """Invokes the optimizer step and convergence criterion."""
          if seed_is_none:
              seed, next_seed = None, previous.seed
          else:
              seed, next_seed = samplers.split_seed(previous.seed, n=2)
          loss, grads, parameters, optimizer_state = optimizer_step_fn(
              previous.parameters, previous.optimizer_state, seed=seed
          )
          if convergence_criterion is None:
              has_converged = previous.has_converged
              convergence_criterion_state = (
                  previous.convergence_criterion_state
              )
          else:
              (
                  has_converged,
                  convergence_criterion_state,
              ) = convergence_criterion.one_step(
                  step,
                  loss,
                  grads,
                  parameters,
                  previous.convergence_criterion_state,
              )
          return MinimizeTraceableQuantities(
              loss=loss,
              gradients=grads,
              parameters=parameters,
              step=step,
              has_converged=has_converged,
              convergence_criterion_state=convergence_criterion_state,
              optimizer_state=optimizer_state,
              seed=next_seed,
          )
  
      def training_loop_body(previous_traced_values, step):
          return ps.cond(
              batch_convergence_reduce_fn(
                  previous_traced_values.has_converged
              ),
              lambda: previous_traced_values,
              lambda: unconverged_training_loop_body(
                  previous_traced_values, step
              ),
          )
  
      return training_loop_body
  
file_path: probability/tensorflow_probability/python/math/minimize.py
imports: import collections
  import functools
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import loop_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.math.gradient import value_and_gradient
label_desc: Use image node_id 3 for calling a global function with example usage: _make_training_loop_body(optimizer_step_fn, convergence_criterion, batch_convergence_reduce_fn, seed_is_none) and returns: training_loop_body, MinimizeTraceableQuantities, ps