code: def testSparseBasic(self):
      self.skipTest("b/195306553")
      if tf.executing_eagerly():
          return
  
      for dtype in [tf.half, tf.float32, tf.float64]:
          with self.cached_session():
              var0 = tf.Variable([[1.1], [2.1]], dtype=dtype)
              var1 = tf.Variable([[3.0], [4.0]], dtype=dtype)
              grads0 = tf.IndexedSlices(
                  tf.constant([0.1], shape=[1, 1], dtype=dtype),
                  tf.constant([0]),
                  tf.constant([2, 1]),
              )
              grads1 = tf.IndexedSlices(
                  tf.constant([0.01], shape=[1, 1], dtype=dtype),
                  tf.constant([1]),
                  tf.constant([2, 1]),
              )
              decay_rate = 0.9
              sgd_op = sgld.StochasticGradientLangevinDynamics(
                  3.0, preconditioner_decay_rate=decay_rate
              ).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
              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]],
                  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 * 0], [4.0 - 3.0 * grads_scaled]],
                  self.evaluate(var1),
              )
  
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 7 for calling the StochasticGradientLangevinDynamicsOptimizerTest obj's underlying member method code with example usage: obj.testSparseBasic() without return types