docstr: Test that SGLD step is computed correctly for a 3D Gaussian energy.
docstr_tok: ["Test","that","SGLD","step","is","computed","correctly","for","a","3D","Gaussian","energy","."]
code: def testPreconditionerComputedCorrectly(self):
      """Test that SGLD step is computed correctly for a 3D Gaussian energy."""
      if tf.executing_eagerly():
          return
  
      with self.cached_session():
          dtype = np.float32
          # Target function is the energy function of normal distribution
          true_mean = dtype([0, 0, 0])
          true_cov = dtype(
              [[1, 0.25, 0.25], [0.25, 1, 0.25], [0.25, 0.25, 1]]
          )
          # Target distribution is defined through the Cholesky decomposition
          chol = tf.linalg.cholesky(true_cov)
          target = mvn_tril.MultivariateNormalTriL(
              loc=true_mean, scale_tril=chol
          )
          var_1 = tf.Variable(name="var_1", initial_value=[1.0, 1.0])
          var_2 = tf.Variable(name="var_2", initial_value=[1.0])
  
          var = [var_1, var_2]
  
          # Set up the learning rate and the optimizer
          learning_rate = 0.5
          optimizer_kernel = sgld.StochasticGradientLangevinDynamics(
              learning_rate=learning_rate, burnin=1
          )
  
          # Target function
          def target_fn(x, y):
              # Stack the input tensors together
              z = tf.concat([x, y], axis=-1) - true_mean
              return -target.log_prob(z)
  
          grads = tf.gradients(ys=target_fn(*var), xs=var)
  
          # Update value of `var` with one iteration of the SGLD (without the
          # normal perturbation, since `burnin > 0`)
          step = optimizer_kernel.apply_gradients(zip(grads, var))
  
          # True theoretical value of `var` after one iteration
          decay_tensor = tf.cast(
              optimizer_kernel._decay_tensor, var[0].dtype
          )
          diagonal_bias = tf.cast(
              optimizer_kernel._diagonal_bias, var[0].dtype
          )
          learning_rate = tf.cast(
              optimizer_kernel._learning_rate, var[0].dtype
          )
          velocity = [
              (
                  decay_tensor * tf.ones_like(v)
                  + (1 - decay_tensor) * tf.square(g)
              )
              for v, g in zip(var, grads)
          ]
          preconditioner = [
              tf.math.rsqrt(vel + diagonal_bias) for vel in velocity
          ]
          # Compute second order gradients
          _, grad_grads = diag_jacobian(xs=var, ys=grads)
          # Compute gradient of the preconditioner (compute the gradient manually)
          preconditioner_grads = [
              -(g * g_g * (1.0 - decay_tensor) * p**3.0)
              for g, g_g, p in zip(grads, grad_grads, preconditioner)
          ]
  
          # True theoretical value of `var` after one iteration
          var_true = [
              v - learning_rate * 0.5 * (p * g - p_g)
              for v, p, g, p_g in zip(
                  var, preconditioner, grads, preconditioner_grads
              )
          ]
          self.evaluate(tf1.global_variables_initializer())
          var_true_ = self.evaluate(var_true)
          self.evaluate(step)
          var_ = self.evaluate(var)  # new `var` after one SGLD step
          self.assertAllClose(var_true_, var_, atol=0.001, rtol=0.001)
  
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 8 for calling the StochasticGradientLangevinDynamicsOptimizerTest obj's underlying member method code with example usage: obj.testPreconditionerComputedCorrectly() and returns: unknown