docstr: Test that for the SGLD finds minimum of the 3D Gaussian energy.
docstr_tok: ["Test","that","for","the","SGLD","finds","minimum","of","the","3D","Gaussian","energy","."]
code: def testDiffusionBehavesCorrectly(self):
      """Test that for the SGLD finds minimum of the 3D Gaussian energy."""
      if tf.executing_eagerly():
          return
  
      with self.cached_session():
          # Set up random seed for the optimizer
          tf.random.set_seed(42)
          dtype = np.float32
          true_mean = dtype([0, 0, 0])
          true_cov = dtype(
              [[1, 0.25, 0.25], [0.25, 1, 0.25], [0.25, 0.25, 1]]
          )
          # Loss is defined through the Cholesky decomposition
          chol = tf.linalg.cholesky(true_cov)
          var_1 = tf.Variable(name="var_1", initial_value=[1.0, 1.0])
          var_2 = tf.Variable(name="var_2", initial_value=[1.0])
  
          # Loss function
          def loss_fn():
              var = tf.concat([var_1, var_2], axis=-1)
              loss_part = tf.linalg.cholesky_solve(
                  chol, var[..., tf.newaxis]
              )
              return tf.linalg.matvec(loss_part, var, transpose_a=True)
  
          # Set up the learning rate with a polynomial decay
          global_step = tf1.train.get_or_create_global_step()
          starter_learning_rate = 0.3
          end_learning_rate = 1e-4
          decay_steps = 1e4
          learning_rate = tf1.train.polynomial_decay(
              starter_learning_rate,
              global_step,
              decay_steps,
              end_learning_rate,
              power=1.0,
          )
  
          # Set up the optimizer
          optimizer_kernel = sgld.StochasticGradientLangevinDynamics(
              learning_rate=learning_rate,
              preconditioner_decay_rate=0.99,
          )
          optimizer_kernel.iterations = global_step
          optimizer = optimizer_kernel.minimize(
              loss_fn, var_list=[var_1, var_2]
          )
  
          # Number of training steps
          training_steps = 5000
          # Record the steps as and treat them as samples
          samples = [
              np.zeros([training_steps, 2]),
              np.zeros([training_steps, 1]),
          ]
          self.evaluate(tf1.global_variables_initializer())
          for step in range(training_steps):
              self.evaluate(optimizer)
              sample = [self.evaluate(var_1), self.evaluate(var_2)]
              samples[0][step, :] = sample[0]
              samples[1][step, :] = sample[1]
  
      samples_ = np.concatenate(samples, axis=-1)
      sample_mean = np.mean(samples_, 0)
      self.assertAllClose(sample_mean, true_mean, atol=0.15, rtol=0.1)
  
file_path: probability/tensorflow_probability/python/optimizer/sgld_test.py
imports: import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import mvn_tril
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math.diag_jacobian import diag_jacobian
  from tensorflow_probability.python.optimizer import sgld
  from tensorflow.python.framework import test_util
label_desc: Use image node_id 9 for calling the StochasticGradientLangevinDynamicsOptimizerTest obj's underlying member method code with example usage: obj.testDiffusionBehavesCorrectly() and returns: tf