code: def testTensorDecayLearningRate(self):
      if not 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.constant([0.1, 0.1], dtype=dtype)
              grads1 = tf.constant([0.01, 0.01], dtype=dtype)
              lrate = tf.Variable(3.0)
              decay_rate = 0.5
              batch_size = 2
              total_num_examples = 10
              optimizer = variational_sgd.VariationalSGD(
                  batch_size,
                  total_num_examples,
                  max_learning_rate=lrate,
                  burnin=0,
                  preconditioner_decay_rate=decay_rate,
              )
              if not tf.executing_eagerly():
                  sgd_op = optimizer.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
              if not tf.executing_eagerly():
                  self.evaluate(sgd_op)
              else:
                  optimizer.apply_gradients(
                      zip([grads0, grads1], [var0, var1])
                  )
              # Validate updated params
              self.assertAllCloseAccordingToType(
                  [1.1 - 3.0 * 0.1, 2.1 - 3.0 * 0.1],
                  self.evaluate(var0),
              )
              self.assertAllCloseAccordingToType(
                  [3.0 - 3.0 * 0.01, 4.0 - 3.0 * 0.01],
                  self.evaluate(var1),
              )
              # Update learning rate to 0
              if not tf.executing_eagerly():
                  self.evaluate(lrate.assign_add(-3.0))
                  self.evaluate(sgd_op)
              else:
                  optimizer._max_learning_rate -= 3.0
                  optimizer.apply_gradients(
                      zip([grads0, grads1], [var0, var1])
                  )
  
              # Validate params haven't changed
              self.assertAllCloseAccordingToType(
                  [1.1 - 3.0 * 0.1, 2.1 - 3.0 * 0.1],
                  self.evaluate(var0),
              )
              self.assertAllCloseAccordingToType(
                  [3.0 - 3.0 * 0.01, 4.0 - 3.0 * 0.01],
                  self.evaluate(var1),
              )
  
              if not tf.executing_eagerly():
                  with self.assertRaises(
                      tf.errors.InvalidArgumentError
                  ):
                      self.evaluate(lrate.assign_add(-3.0))
                      self.evaluate(sgd_op)
  
file_path: probability/tensorflow_probability/python/optimizer/variational_sgd_test.py
imports: import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.optimizer import variational_sgd
  from tensorflow.python.framework import test_util
label_desc: Use image node_id 5 for calling the VariationalSGDTest obj's underlying member method code with example usage: obj.testTensorDecayLearningRate() without return types