code: def testWithGlobalStep(self):
      if tf.executing_eagerly():
          return
  
      for dtype in [tf.float32, tf.float64]:
          with self.cached_session():
              step = tf.Variable(0, dtype=tf.int64)
  
              var0 = tf.Variable([1.1, 2.1], dtype=dtype)
              var1 = tf.Variable([3.0, 4.0], dtype=dtype)
              grads0 = tf.constant([0.1, 0.1], dtype=dtype)
              grads1 = tf.constant([0.01, 0.01], dtype=dtype)
              decay_rate = 0.1
  
              sgd_opt = sgld.StochasticGradientLangevinDynamics(
                  3.0, preconditioner_decay_rate=decay_rate
              )
              sgd_opt.iterations = step
              sgd_op = sgd_opt.apply_gradients(
                  zip([grads0, grads1], [var0, var1])
              )
  
              self.evaluate(tf1.global_variables_initializer())
  
              # Fetch params to validate initial values
              self.assertAllCloseAccordingToType(
                  [1.1, 2.1], self.evaluate(var0)
              )
              self.assertAllCloseAccordingToType(
                  [3.0, 4.0], self.evaluate(var1)
              )
              # Run 1 step of sgd
              self.evaluate(sgd_op)
  
              # Validate updated params and step
              grads_scaled = (
                  0.5
                  * 0.1
                  / np.sqrt(
                      decay_rate + (1.0 - decay_rate) * 0.1**2 + 1e-8
                  )
              )
              # Note that `tfp.math.diag_jacobian(xs=var, ys=grad)` returns zero
              # tensor
              self.assertAllCloseAccordingToType(
                  [1.1 - 3.0 * grads_scaled, 2.1 - 3.0 * grads_scaled],
                  self.evaluate(var0),
              )
              grads_scaled = (
                  0.5
                  * 0.01
                  / np.sqrt(
                      decay_rate + (1.0 - decay_rate) * 0.01**2 + 1e-8
                  )
              )
              self.assertAllCloseAccordingToType(
                  [3.0 - 3.0 * grads_scaled, 4.0 - 3.0 * grads_scaled],
                  self.evaluate(var1),
              )
              self.assertAllCloseAccordingToType(1, self.evaluate(step))
  
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 6 for calling the StochasticGradientLangevinDynamicsOptimizerTest obj's underlying member method code with example usage: obj.testWithGlobalStep() without return types