code: def test_stochastic_optimization(self):
      seed = test_util.test_seed_stream()
      tf.random.set_seed(seed())
  
      # Example of fitting one normal to another using a
      # Monte Carlo variational loss.
      locs = tf.Variable(tf.random.normal([10], seed=seed()))
      scales = deferred_tensor.TransformedVariable(
          tf.nn.softplus(tf.random.normal([10], seed=seed())),
          softplus.Softplus(),
      )
      trained_dist = normal.Normal(locs, scales)
      target_dist = normal.Normal(loc=-0.4, scale=1.2)
  
      optimizer = tf_keras.optimizers.Adam(learning_rate=0.1)
  
      @tf.function(autograph=False)
      def optimization_step():
          with tf.GradientTape() as tape:
              loss = csiszar_divergence.monte_carlo_variational_loss(
                  target_log_prob_fn=target_dist.log_prob,
                  surrogate_posterior=trained_dist,
                  sample_size=20,
                  seed=seed(),
              )
          grads = tape.gradient(loss, trained_dist.trainable_variables)
          optimizer.apply_gradients(
              zip(grads, trained_dist.trainable_variables)
          )
          return loss, grads
  
      criterion = sgau.SuccessiveGradientsAreUncorrelated(
          window_size=10, min_num_steps=20
      )
  
      loss, grads = optimization_step()
      self.evaluate(tf1.global_variables_initializer())
      auxiliary_state = criterion.bootstrap(
          loss, grads, trained_dist.trainable_variables
      )
      for step in range(1, 100):
          loss, grads = optimization_step()
          has_converged, auxiliary_state = criterion.one_step(
              step,
              loss,
              grads,
              trained_dist.trainable_variables,
              auxiliary_state,
          )
  
          has_converged_ = self.evaluate(has_converged)
          if has_converged_:
              break
  
      # Check that the criterion successfully stopped the optimization
      # (at step 32 with the test seed as of this writing).
      self.assertLess(
          step, 99
      )  # pylint: disable=undefined-loop-variable
  
      # Because this is a stochastic optimization with no learning rate decay,
      # we will not converge to the true values, just to a stationary distribution
      # that (hopefully) includes them.
      self.assertLess(self.evaluate(tf.reduce_sum(loss)), 1.5)
      self.assertAllClose(
          *self.evaluate((tf.reduce_mean(locs), target_dist.mean())),
          atol=0.5,
      )
      self.assertAllClose(
          *self.evaluate(
              (tf.reduce_mean(scales), target_dist.stddev())
          ),
          atol=0.5,
      )
  
file_path: probability/tensorflow_probability/python/optimizer/convergence_criteria/successive_gradients_are_uncorrelated_test.py
imports: import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import softplus
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.internal import tf_keras
  from tensorflow_probability.python.optimizer.convergence_criteria import successive_gradients_are_uncorrelated
  from tensorflow_probability.python.util import deferred_tensor
  from tensorflow_probability.python.vi import csiszar_divergence
label_desc: Use image node_id 1 for calling the SuccessiveGradientsAreUncorrelatedTests obj's underlying member method code with example usage: obj.test_stochastic_optimization() and returns: loss,  grads