class_docstr: An optimizer module for constant stochastic gradient descent.
  
  This implements an optimizer module for the constant stochastic gradient
  descent algorithm [(Mandt et al., 2017)][1]. The optimization variable is
  regarded as an approximate sample from the posterior .
  
  Note: If a prior is included in the loss, it should be scaled by
  `1/num_pseudo_batches`, where num_pseudo_batches is the number of minibatches
  in the data.  I.e., it should be divided by the `num_pseudo_batches` term
  described below.
  
  Args:
    batch_size: Scalar `int`-like `Tensor`. The number of examples in a
      minibatch in the data set. Note: Assumes the loss is taken as the mean
      over a minibatch. Otherwise if the sum was taken set this to 1.
    total_num_examples: Scalar `int`-like `Tensor`. The total number of examples
      in the data set.
    max_learning_rate: Scalar `float`-like `Tensor`. A maximum allowable
      effective coordinate-wise learning rate. The algorithm scales down any
      effective learning rate (i.e. after preconditioning) that is larger than
      this. (Default: `1`)
    preconditioner_decay_rate: Scalar `float`-like `Tensor`. The exponential
      decay rate of the rescaling of the preconditioner (RMSprop). (This is
      "alpha" in Mandt et al. (2017)). Should be smaller than but nearly `1` to
      approximate sampling from the posterior. (Default: `0.95`)
    burnin: Scalar `int`-like `Tensor`. The number of iterations to collect
      gradient statistics to update the preconditioner before starting to draw
      noisy samples. (Default: `25`)
    burnin_max_learning_rate: Scalar `float`-like `Tensor`. Maximum learning
      rate to use during the burnin period.
      (Default: `1e-8`)
    use_single_learning_rate: Boolean Indicates whether one single learning
      rate is used or coordinate_wise learning rates are used.
      (Default: `False`)
    name: Python `str` describing ops managed by this function.
      (Default: `"VariationalSGD"`)
  
  Raises:
    InvalidArgumentError: If preconditioner_decay_rate is a `Tensor` not in
      `(0,1]`.
  
  #### References
  
  [1]: Stephan Mandt, Matthew D. Hoffman, and David M. Blei. Stochastic
       Gradient Descent as Approximate Bayesian Inference. _arXiv preprint
       arXiv:1704.04289_, 2017. https://arxiv.org/abs/1704.04289
code: def __init__(
      self,
      batch_size,
      total_num_examples,
      max_learning_rate=1.0,
      preconditioner_decay_rate=0.95,
      burnin=25,
      burnin_max_learning_rate=1e-6,
      use_single_learning_rate=False,
      name=None,
  ):
      default_name = "VariationalSGD"
      with tf.name_scope(name or default_name):
          self._preconditioner_decay_rate = tf.convert_to_tensor(
              preconditioner_decay_rate,
              name="preconditioner_decay_rate",
          )
          self._batch_size = tf.convert_to_tensor(
              batch_size, name="batch_size"
          )
          self._total_num_examples = tf.convert_to_tensor(
              total_num_examples, name="total_num_examples"
          )
  
          self._burnin = tf.convert_to_tensor(
              burnin,
              name="burnin",
              dtype=dtype_util.common_dtype(
                  [burnin], dtype_hint=tf.int64
              ),
          )
          self._burnin_max_learning_rate = tf.convert_to_tensor(
              burnin_max_learning_rate, name="burnin_max_learning_rate"
          )
          self._max_learning_rate = tf.convert_to_tensor(
              max_learning_rate, name="max_learning_rate"
          )
          self._use_single_learning_rate = use_single_learning_rate
  
          self._preconditioner_decay_rate = distribution_util.with_dependencies(
              [
                  assert_util.assert_non_negative(
                      self._preconditioner_decay_rate,
                      message="`preconditioner_decay_rate` must be non-negative",
                  ),
                  assert_util.assert_less_equal(
                      self._preconditioner_decay_rate,
                      1.0,
                      message="`preconditioner_decay_rate` must be at most 1.",
                  ),
              ],
              self._preconditioner_decay_rate,
          )
  
          self._batch_size = distribution_util.with_dependencies(
              [
                  assert_util.assert_greater(
                      self._batch_size,
                      0,
                      message="`batch_size` must be greater than zero",
                  )
              ],
              self._batch_size,
          )
  
          self._total_num_examples = distribution_util.with_dependencies(
              [
                  assert_util.assert_greater(
                      self._total_num_examples,
                      0,
                      message="`total_num_examples` must be greater than zero",
                  )
              ],
              self._total_num_examples,
          )
  
          self._burnin = distribution_util.with_dependencies(
              [
                  assert_util.assert_non_negative(
                      self._burnin,
                      message="`burnin` must be non-negative",
                  ),
                  assert_util.assert_integer(
                      self._burnin,
                      message="`burnin` must be an integer",
                  ),
              ],
              self._burnin,
          )
  
          self._burnin_max_learning_rate = distribution_util.with_dependencies(
              [
                  assert_util.assert_non_negative(
                      self._burnin_max_learning_rate,
                      message="`burnin_max_learning_rate` must be non-negative",
                  )
              ],
              self._burnin_max_learning_rate,
          )
  
          self._max_learning_rate = distribution_util.with_dependencies(
              [
                  assert_util.assert_non_negative(
                      self._max_learning_rate,
                      message="`max_learning_rate` must be non-negative",
                  )
              ],
              self._max_learning_rate,
          )
  
          super(VariationalSGD, self).__init__(
              name=name or default_name
          )
  
file_path: probability/tensorflow_probability/python/optimizer/variational_sgd.py
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import tf_keras
label_desc: Use image node_id 1 to create a new VariationalSGD object from inherited base classes: tf_keras.optimizers.legacy with example: obj = VariationalSGD(batch_size, total_num_examples, max_learning_rate, preconditioner_decay_rate, burnin, burnin_max_learning_rate, use_single_learning_rate, name)