code: def test_variational_em(self):
      seed = test_util.test_seed()
  
      num_samples = 10000
      mu, sigma = 3.0, 5.0
      x = test_util.test_np_rng().randn(num_samples) * sigma + mu
  
      # Test that the tape automatically picks up any trainable variables in
      # the model, even though it's just a function with no explicit
      # `.trainable_variables`
      likelihood_scale = deferred_tensor.TransformedVariable(
          1.0, softplus.Softplus(), name="scale"
      )
  
      def trainable_log_prob(z):
          lp = normal.Normal(0.0, 1.0).log_prob(z)
          lp += tf.reduce_sum(
              normal.Normal(
                  z[..., tf.newaxis], likelihood_scale
              ).log_prob(x),
              axis=-1,
          )
          return lp
  
      # For this simple normal-normal model, the true posterior is also normal.
      z_posterior_precision = 1.0 / sigma**2 * num_samples + 1.0**2
      z_posterior_stddev = np.sqrt(1.0 / z_posterior_precision)
      z_posterior_mean = (
          1.0 / sigma**2 * num_samples * mu
      ) / z_posterior_precision
  
      q_loc = tf.Variable(0.0, name="mu")
      q_scale = deferred_tensor.TransformedVariable(
          1.0, softplus.Softplus(), name="q_scale"
      )
      q = normal.Normal(q_loc, q_scale)
      loss_curve = optimization.fit_surrogate_posterior(
          trainable_log_prob,
          q,
          num_steps=1000,
          sample_size=10,
          optimizer=tf_keras.optimizers.Adam(0.1),
          seed=seed,
      )
      self.evaluate(tf1.global_variables_initializer())
      with tf.control_dependencies([loss_curve]):
          final_q_loc = tf.identity(q.mean())
          final_q_scale = tf.identity(q.stddev())
          final_likelihood_scale = tf.convert_to_tensor(
              likelihood_scale
          )
  
      # We expect to recover the true posterior because the variational family
      # includes the true posterior, and the true parameters because we observed
      # a large number of sampled points.
      (
          final_likelihood_scale_,
          final_q_loc_,
          final_q_scale_,
      ) = self.evaluate(
          (final_likelihood_scale, final_q_loc, final_q_scale)
      )
      self.assertAllClose(final_likelihood_scale_, sigma, atol=0.2)
      self.assertAllClose(final_q_loc_, z_posterior_mean, atol=0.2)
      self.assertAllClose(final_q_scale_, z_posterior_stddev, atol=0.1)
  
imports: import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import fill_scale_tril
  from tensorflow_probability.python.bijectors import softplus
  from tensorflow_probability.python.distributions import deterministic
  from tensorflow_probability.python.distributions import gaussian_process
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import joint_distribution_auto_batched
  from tensorflow_probability.python.distributions import joint_distribution_coroutine
  from tensorflow_probability.python.distributions import lognormal
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import poisson
  from tensorflow_probability.python.experimental.util import trainable
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.internal import tf_keras
  from tensorflow_probability.python.math.psd_kernels import exponentiated_quadratic
  from tensorflow_probability.python.util import deferred_tensor
  from tensorflow_probability.python.vi import optimization
label_desc: Use image node_id 1 for calling the OptimizationTests obj's underlying member method code with example usage: obj.test_variational_em() and returns: lp