code: def test_custom_eager_optimization_loop(self):
      if not tf.executing_eagerly():
          return
  
      batch_shape = [2, 3]
      num_timesteps = 5
      observed_time_series = self._build_tensor(
          np.random.randn(*(batch_shape + [num_timesteps]))
      )
  
      model = self._build_model(observed_time_series)
  
      surrogate_posterior = fitting.build_factored_surrogate_posterior(
          model=model
      )
      self.assertLen(
          surrogate_posterior.trainable_variables,
          len(model.parameters) * 2,
      )  # Loc and scale for each param.
  
      @tf.function(
          autograph=False
      )  # Ensure the loss is computed efficiently
      def loss_fn(sample_size=3):
          return csiszar_divergence.monte_carlo_variational_loss(
              model.joint_distribution(observed_time_series).log_prob,
              surrogate_posterior,
              sample_size=sample_size,
          )
  
      initial_loss = self.evaluate(loss_fn(sample_size=10))
  
      # TODO(b/137299119) Replace with TF2 optimizer.
      optimizer = tf1.train.AdamOptimizer(learning_rate=0.1)
      for _ in range(10):
          with tf.GradientTape() as tape:
              loss = loss_fn()
          grads = tape.gradient(
              loss, surrogate_posterior.trainable_variables
          )
          self.evaluate(
              optimizer.apply_gradients(
                  zip(grads, surrogate_posterior.trainable_variables)
              )
          )
  
      final_loss = self.evaluate(loss_fn(sample_size=10))
      self.assertAllEqual(final_loss.shape, batch_shape)
      self.assertLess(np.mean(final_loss), np.mean(initial_loss))
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.sts import fitting
  from tensorflow_probability.python.sts.components import local_linear_trend
  from tensorflow_probability.python.sts.components import seasonal
  from tensorflow_probability.python.sts.components import sum
  from tensorflow_probability.python.sts.internal import missing_values_util
  from tensorflow_probability.python.vi import csiszar_divergence
  from tensorflow_probability.python.vi import optimization
label_desc: Use image node_id 4 for calling the _VariationalInferenceTests obj's underlying member method code with example usage: obj.test_custom_eager_optimization_loop() and returns: csiszar_divergence