code: def _make_stateless_optimizer_step_fn(loss_fn, optimizer):
      """Constructs a single step of a pure functional optimizer."""
  
      @tf.function(autograph=False)
      def optimizer_step(parameters, optimizer_state, seed=None):
          """Runs a single optimization step."""
          try:
              loss, grads = value_and_gradient(
                  functools.partial(loss_fn, seed=seed), parameters
              )
          except TypeError:
              loss, grads = value_and_gradient(loss_fn, parameters)
          # Coerce grads to the same sequence type (e.g., namedtuple) as parameters.
          grads = tf.nest.pack_sequence_as(
              parameters, tf.nest.flatten(grads)
          )
          updates, optimizer_state = optimizer.update(
              grads, optimizer_state, parameters
          )
          # Apply updates.
          parameters = tf.nest.map_structure(
              lambda a, b: a + b, parameters, updates
          )
          return loss, grads, parameters, optimizer_state
  
      return optimizer_step
  
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 5 for calling a global function with example usage: _make_stateless_optimizer_step_fn(loss_fn, optimizer) and returns: optimizer_step, loss,  grads,  parameters,  optimizer_state