code: def _minimize_common(
      num_steps,
      optimizer_step_fn,
      initial_parameters,
      initial_optimizer_state,
      convergence_criterion,
      batch_convergence_reduce_fn,
      trace_fn,
      return_full_length_trace=True,
      jit_compile=False,
      seed=None,
      name="minimize",
  ):
      """General-purpose optimization loop."""
      if jit_compile:
          # Run the entire minimization inside a jit-compiled function. This is
          # typically faster than jit-compiling the individual steps.
          kwargs = dict(locals())
          kwargs["jit_compile"] = False
  
          @tf.function(autograph=False, jit_compile=True)
          def run_jitted_minimize():
              return _minimize_common(**kwargs)
  
          return run_jitted_minimize()
  
      # Main optimization routine.
      with tf.name_scope(name) as name:
          seed_is_none = seed is None
          if not seed_is_none:
              seed = samplers.sanitize_seed(seed, salt="minimize")
  
          if not return_full_length_trace:
              # Augment trace to record convergence info, so we can truncate it later.
              trace_fn = _trace_has_converged(
                  trace_fn, batch_convergence_reduce_fn
              )
  
          # Take an initial training step to obtain the initial loss and values, which
          # are used to initialize the convergence criterion. This will trigger
          # tf.function tracing of `optimizer_step_fn`, which is
          # then reused inside the training loop (i.e., it is only traced once).
          (
              initial_loss,
              initial_grads,
              initial_parameters,
              initial_optimizer_state,
          ) = optimizer_step_fn(
              parameters=initial_parameters,
              optimizer_state=initial_optimizer_state,
              seed=seed,
          )
  
          initial_convergence_criterion_state = ()
          if convergence_criterion is not None:
              initial_convergence_criterion_state = (
                  convergence_criterion.bootstrap(
                      initial_loss, initial_grads, initial_parameters
                  )
              )
  
          with tf.control_dependencies([initial_loss]):
              has_converged = tf.zeros(
                  ps.shape(initial_loss), dtype=tf.bool
              )
              initial_state = MinimizeTraceableQuantities(
                  loss=initial_loss,
                  gradients=initial_grads,
                  parameters=initial_parameters,
                  step=0,
                  has_converged=has_converged,
                  convergence_criterion_state=initial_convergence_criterion_state,
                  optimizer_state=initial_optimizer_state,
                  seed=-1 if seed_is_none else seed,
              )
              initial_traced_values = tf.nest.map_structure(
                  lambda x: tf.convert_to_tensor(x)[tf.newaxis, ...],
                  trace_fn(initial_state),
              )
  
          def run_optimization_loop():
              (
                  final_step_traceable_values,
                  loop_traced_values,
              ) = loop_util.trace_scan(
                  loop_fn=_make_training_loop_body(
                      optimizer_step_fn=optimizer_step_fn,
                      convergence_criterion=convergence_criterion,
                      batch_convergence_reduce_fn=batch_convergence_reduce_fn,
                      seed_is_none=seed_is_none,
                  ),
                  initial_state=initial_state,
                  elems=tf.range(1, num_steps),
                  trace_fn=trace_fn,
              )
              traced_values = tf.nest.map_structure(
                  lambda t0, t: tf.concat([t0, t], axis=0),
                  initial_traced_values,
                  loop_traced_values,
              )
              return (
                  final_step_traceable_values.parameters,
                  traced_values,
              )
  
          # When variables are involved, we want to make sure the initial trace is
          # sequenced before the rest of the optimization loop. Otherwise, it's
          # possible for the loop to *complete* before initial_trace_values are
          # executed, causing nonsensical traces.
          with tf.control_dependencies(
              tf.nest.flatten(initial_traced_values)
          ):
              final_parameters, traced_values = ps.cond(
                  num_steps > 1,
                  run_optimization_loop,
                  lambda: (initial_parameters, initial_traced_values),
              )
  
          if not return_full_length_trace:
              traced_values = _truncate_at_has_converged(traced_values)
          return final_parameters, traced_values
  
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 4 for calling a global function with example usage: _minimize_common(num_steps, optimizer_step_fn, initial_parameters, initial_optimizer_state, convergence_criterion, batch_convergence_reduce_fn, trace_fn, return_full_length_trace, jit_compile, seed, name) and returns: run_jitted_minimize, final_parameters,  traced_values, _minimize_common, final_step_traceable_values,  traced_values