code: def testBurnin(self):
      for burnin_dtype in [tf.int8, tf.int16, tf.int32, tf.int64]:
          with self.cached_session():
              var0 = tf.Variable([1.1, 2.1], dtype=tf.float32)
              grads0 = tf.constant([0.1, 0.1], dtype=tf.float32)
              decay_rate = 0.53
              sgd_optimizer = variational_sgd.VariationalSGD(
                  1,
                  1,
                  preconditioner_decay_rate=decay_rate,
                  burnin=tf.constant(10, dtype=burnin_dtype),
                  max_learning_rate=3.0,
                  burnin_max_learning_rate=3.0,
                  use_single_learning_rate=True,
              )
  
              if not tf.executing_eagerly():
                  sgd_op = sgd_optimizer.apply_gradients(
                      [(grads0, var0)]
                  )
  
              self.evaluate(tf1.global_variables_initializer())
  
              # Validate iterations is initialized to 0.
              self.assertAllCloseAccordingToType(
                  0, self.evaluate(sgd_optimizer.iterations)
              )
              # Run 1 step of sgd
              if not tf.executing_eagerly():
                  self.evaluate(sgd_op)
              else:
                  sgd_optimizer.apply_gradients([(grads0, var0)])
  
              # Validate iterations is incremented.
              self.assertAllCloseAccordingToType(
                  1, self.evaluate(sgd_optimizer.iterations)
              )
  
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 4 for calling the VariationalSGDTest obj's underlying member method code with example usage: obj.testBurnin() without return types