docstr: Check that the gradient of the volatility is computed correctly.
docstr_tok: ["Check","that","the","gradient","of","the","volatility","is","computed","correctly","."]
code: def testLangevinCorrectVolatilityGradient(self):
      """Check that the gradient of the volatility is computed correctly."""
      # Consider the example target distribution as in `testLangevin3DNormal`
      dtype = np.float32
      true_mean = dtype([1, 2, 7])
      true_cov = dtype(
          [[1, 0.25, 0.25], [0.25, 1, 0.25], [0.25, 0.25, 1]]
      )
      num_chains = 100
  
      chol = tf.linalg.cholesky(true_cov)
      target = mvn_tril.MultivariateNormalTriL(
          loc=true_mean, scale_tril=chol
      )
  
      def target_log_prob(x, y):
          # Stack the input tensors together
          z = tf.concat([x, y], axis=-1)
          return target.log_prob(z)
  
      def volatility_fn(x, y):
          # Stack the input tensors together
          return [
              1.0 / (0.5 + 0.1 * tf.abs(x + y)),
              1.0 / (0.5 + 0.1 * tf.abs(y)),
          ]
  
      # Initial state of the chain
      init_state = [
          np.ones([num_chains, 2], dtype=dtype),
          np.ones([num_chains, 1], dtype=dtype),
      ]
  
      # Define MALA with constant volatility
      langevin_unit = langevin.MetropolisAdjustedLangevinAlgorithm(
          target_log_prob_fn=target_log_prob, step_size=0.1
      )
      # Define MALA with volatility being `volatility_fn`
      langevin_general = langevin.MetropolisAdjustedLangevinAlgorithm(
          target_log_prob_fn=target_log_prob,
          step_size=0.1,
          volatility_fn=volatility_fn,
      )
  
      # Initialize the samplers
      kernel_unit_volatility = langevin_unit.bootstrap_results(
          init_state
      )
      kernel_general = langevin_general.bootstrap_results(init_state)
  
      # For `langevin_general` volatility gradient should be zero.
      (
          grad_1,
          grad_2,
      ) = kernel_unit_volatility.accepted_results.grads_volatility
      self.assertAllEqual(
          self.evaluate(grad_1),
          np.zeros(shape=init_state[0].shape, dtype=dtype),
      )
      self.assertAllEqual(
          self.evaluate(grad_2),
          np.zeros(shape=init_state[1].shape, dtype=dtype),
      )
  
      # For `langevin_unit` volatility gradient should be around -0.926 for
      # each direction.
      grad_1, grad_2 = kernel_general.accepted_results.grads_volatility
      self.assertAllClose(
          self.evaluate(grad_1),
          -0.583 * np.ones(shape=init_state[0].shape, dtype=dtype),
          atol=0.01,
          rtol=0.01,
      )
      self.assertAllClose(
          self.evaluate(grad_2),
          -0.926 * np.ones(shape=init_state[1].shape, dtype=dtype),
          atol=0.01,
          rtol=0.01,
      )
  
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import mvn_tril
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.internal import distribute_lib
  from tensorflow_probability.python.internal import distribute_test_lib
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.mcmc import langevin
  from tensorflow_probability.python.mcmc import sample
label_desc: Use image node_id 4 for calling the LangevinTest obj's underlying member method code with example usage: obj.testLangevinCorrectVolatilityGradient() and returns: target, list