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 _get_coordinatewise_learning_rate(self, grad, var):
      # Compute the learning rate using a moving average for the diagonal of BB^T
      avg_first = self.get_slot(var, "first_moment")
      avg_second = self.get_slot(var, "second_moment")
      decay_tensor = tf.cast(self._decay_tensor, var.dtype)
      batch_size = tf.cast(self._batch_size_tensor, var.dtype)
  
      # Create an estimator for the moving average of gradient mean and variance
      # via Welford's algorithm
      if isinstance(grad, tf.Tensor):
          delta = grad - avg_first
          first_moment_update = avg_first.assign_add(
              delta
              * tf.where(
                  self.iterations < 1,
                  dtype_util.as_numpy_dtype(var.dtype)(1.0),
                  1.0 - decay_tensor,
              )
          )
  
          with tf.control_dependencies([first_moment_update]):
              second_moment_update = avg_second.assign_add(
                  tf.cast(self.iterations < 1, var.dtype)
                  * -(1.0 - decay_tensor)
                  * (avg_second - decay_tensor * tf.square(delta))
              )
          diag_preconditioner = distribution_util.with_dependencies(
              [second_moment_update],
              tf.clip_by_value(avg_second, 1e-12, 1e12),
          )
      elif isinstance(grad, tf.IndexedSlices):
          delta = grad.values - tf.gather_nd(avg_first, grad.indices)
          first_moment_update = tf.compat.v1.scatter_add(
              avg_first,
              grad.indices,
              delta
              * tf.where(
                  self.iterations < 1,
                  dtype_util.as_numpy_dtype(var.dtype)(1.0),
                  1.0 - decay_tensor,
              ),
          )
  
          with tf.control_dependencies([first_moment_update]):
              avg_second = tf.compat.v1.scatter_add(
                  avg_second,
                  grad.indices,
                  tf.cast(self.iterations < 1, var.dtype)
                  * -(1.0 - decay_tensor)
                  * (
                      tf.gather_nd(avg_second, grad.indices)
                      - decay_tensor * tf.square(delta)
                  ),
              )
              avg_second = tf.gather_nd(avg_second, grad.indices)
              # TODO(b/70783772): Needs dtype specific clipping.
              diag_preconditioner = tf.clip_by_value(
                  avg_second, 1e-12, 1e12
              )
      else:
          raise tf.errors.InvalidArgumentError(
              None, None, "grad must of type Tensor or IndexedSlice"
          )
  
      diag_preconditioner *= batch_size
  
      if self._use_single_learning_rate:
          diag_preconditioner = tf.reduce_mean(diag_preconditioner)
  
      # From Theorem 2 Corollary 1 of Mandt et al. 2017
      return (
          2.0
          * batch_size
          / (
              tf.cast(self._total_num_examples, var.dtype.base_dtype)
              * diag_preconditioner
          )
      )
  
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 5 for calling the VariationalSGD obj's underlying member method code with example usage: obj._get_coordinatewise_learning_rate(grad, var) and returns: unknown