class_docstr: An optimizer module for stochastic gradient Langevin dynamics.
  
  This implements the preconditioned Stochastic Gradient Langevin Dynamics
  optimizer [(Li et al., 2016)][1]. The optimization variable is regarded as a
  sample from the posterior under Stochastic Gradient Langevin Dynamics with
  noise rescaled in each dimension according to [RMSProp](
  http://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf).
  
  Note: If a prior is included in the loss, it should be scaled by
  `1/data_size`, where `data_size` is the number of points in the data set.
  I.e., it should be divided by the `data_size` term described below.
  
  #### Examples
  
  ##### Optimizing energy of a 3D-Gaussian distribution
  
  This example demonstrates that for a fixed step size SGLD works as an
  approximate version of MALA (tfp.mcmc.MetropolisAdjustedLangevinAlgorithm).
  
  ```python
  import tensorflow as tf
  import tensorflow_probability as tfp
  import numpy as np
  
  tfd = tfp.distributions
  dtype = np.float32
  
  with tf.Session(graph=tf.Graph()) as sess:
    # Set up random seed for the optimizer
    tf.random.set_seed(42)
    true_mean = dtype([0, 0, 0])
    true_cov = dtype([[1, 0.25, 0.25], [0.25, 1, 0.25], [0.25, 0.25, 1]])
    # Loss is defined through the Cholesky decomposition
    chol = tf.linalg.cholesky(true_cov)
  
    var_1 = tf.Variable(name='var_1', initial_value=[1., 1.])
    var_2 = tf.Variable(name='var_2', initial_value=[1.])
  
    def loss_fn():
      var = tf.concat([var_1, var_2], axis=-1)
      loss_part = tf.linalg.cholesky_solve(chol, var[..., tf.newaxis])
      return tf.linalg.matvec(loss_part, var, transpose_a=True)
  
    # Set up the learning rate with a polynomial decay
    step = tf.Variable(0, dtype=tf.int64)
    starter_learning_rate = .3
    end_learning_rate = 1e-4
    decay_steps = 1e4
    learning_rate = tf.compat.v1.train.polynomial_decay(
        starter_learning_rate,
        step,
        decay_steps,
        end_learning_rate,
        power=1.)
  
    # Set up the optimizer
    optimizer_kernel = tfp.optimizer.StochasticGradientLangevinDynamics(
        learning_rate=learning_rate, preconditioner_decay_rate=0.99)
    optimizer_kernel.iterations = step
    optimizer = optimizer_kernel.minimize(loss_fn, var_list=[var_1, var_2])
  
    # Number of training steps
    training_steps = 5000
    # Record the steps as and treat them as samples
    samples = [np.zeros([training_steps, 2]), np.zeros([training_steps, 1])]
    sess.run(tf.compat.v1.global_variables_initializer())
    for step in range(training_steps):
      sess.run(optimizer)
      sample = [sess.run(var_1), sess.run(var_2)]
      samples[0][step, :] = sample[0]
      samples[1][step, :] = sample[1]
  
    samples_ = np.concatenate(samples, axis=-1)
    sample_mean = np.mean(samples_, 0)
    print('sample mean', sample_mean)
  ```
  Args:
    learning_rate: Scalar `float`-like `Tensor`. The base learning rate for the
      optimizer. Must be tuned to the specific function being minimized.
    preconditioner_decay_rate: Scalar `float`-like `Tensor`. The exponential
      decay rate of the rescaling of the preconditioner (RMSprop). (This is
      "alpha" in Li et al. (2016)). Should be smaller than but nearly `1` to
      approximate sampling from the posterior. (Default: `0.95`)
    data_size: Scalar `int`-like `Tensor`. The effective number of
      points in the data set. Assumes that the loss is taken as the mean over a
      minibatch. Otherwise if the sum was taken, divide this number by the
      batch size. If a prior is included in the loss function, it should be
      normalized by `data_size`. Default value: `1`.
    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`)
    diagonal_bias: Scalar `float`-like `Tensor`. Term added to the diagonal of
      the preconditioner to prevent the preconditioner from degenerating.
      (Default: `1e-8`)
    name: Python `str` describing ops managed by this function.
      (Default: `"StochasticGradientLangevinDynamics"`)
    parallel_iterations: the number of coordinates for which the gradients of
        the preconditioning matrix can be computed in parallel. Must be a
        positive integer.
  
  Raises:
    InvalidArgumentError: If preconditioner_decay_rate is a `Tensor` not in
      `(0,1]`.
    NotImplementedError: If eager execution is enabled.
  
  #### References
  
  [1]: Chunyuan Li, Changyou Chen, David Carlson, and Lawrence Carin.
       Preconditioned Stochastic Gradient Langevin Dynamics for Deep Neural
       Networks. In _Association for the Advancement of Artificial
       Intelligence_, 2016. https://arxiv.org/abs/1512.07666
class_docstr_tok: ["An","optimizer","module","for","stochastic","gradient","Langevin","dynamics",".","This","implements","the","preconditioned","Stochastic","Gradient","Langevin","Dynamics","optimizer","[","(","Li","et","al.",",","2016",")","]","[","1","]",".","The","optimization","variable","is","regarded","as","a","sample","from","the","posterior","under","Stochastic","Gradient","Langevin","Dynamics","with","noise","rescaled","in","each","dimension","according","to","[","RMSProp","]","(","http",":","\/\/www.cs.toronto.edu\/~tijmen\/csc321\/slides\/lecture_slides_lec6.pdf",")",".","Note",":","If","a","prior","is","included","in","the","loss",",","it","should","be","scaled","by","`","1\/data_size","`",",","where","`","data_size","`","is","the","number","of","points","in","the","data","set",".","I.e.",",","it","should","be","divided","by","the","`","data_size","`","term","described","below",".","#","#","#","#","Examples","#","#","#","#","#","Optimizing","energy","of","a","3D-Gaussian","distribution","This","example","demonstrates","that","for","a","fixed","step","size","SGLD","works","as","an","approximate","version","of","MALA","(","tfp.mcmc.MetropolisAdjustedLangevinAlgorithm",")",".","``","`","python","import","tensorflow","as","tf","import","tensorflow_probability","as","tfp","import","numpy","as","np","tfd","=","tfp.distributions","dtype","=","np.float32","with","tf.Session","(","graph=tf.Graph","(",")",")","as","sess",":","#","Set","up","random","seed","for","the","optimizer","tf.random.set_seed","(","42",")","true_mean","=","dtype","(","[","0",",","0",",","0","]",")","true_cov","=","dtype","(","[","[","1",",","0.25",",","0.25","]",",","[","0.25",",","1",",","0.25","]",",","[","0.25",",","0.25",",","1","]","]",")","#","Loss","is","defined","through","the","Cholesky","decomposition","chol","=","tf.linalg.cholesky","(","true_cov",")","var_1","=","tf.Variable","(","name='var_1","'",",","initial_value=","[","1.",",","1",".","]",")","var_2","=","tf.Variable","(","name='var_2","'",",","initial_value=","[","1",".","]",")","def","loss_fn","(",")",":","var","=","tf.concat","(","[","var_1",",","var_2","]",",","axis=-1",")","loss_part","=","tf.linalg.cholesky_solve","(","chol",",","var","[","...",",","tf.newaxis","]",")","return","tf.linalg.matvec","(","loss_part",",","var",",","transpose_a=True",")","#","Set","up","the","learning","rate","with","a","polynomial","decay","step","=","tf.Variable","(","0",",","dtype=tf.int64",")","starter_learning_rate","=",".3","end_learning_rate","=","1e-4","decay_steps","=","1e4","learning_rate","=","tf.compat.v1.train.polynomial_decay","(","starter_learning_rate",",","step",",","decay_steps",",","end_learning_rate",",","power=1",".",")","#","Set","up","the","optimizer","optimizer_kernel","=","tfp.optimizer.StochasticGradientLangevinDynamics","(","learning_rate=learning_rate",",","preconditioner_decay_rate=0.99",")","optimizer_kernel.iterations","=","step","optimizer","=","optimizer_kernel.minimize","(","loss_fn",",","var_list=","[","var_1",",","var_2","]",")","#","Number","of","training","steps","training_steps","=","5000","#","Record","the","steps","as","and","treat","them","as","samples","samples","=","[","np.zeros","(","[","training_steps",",","2","]",")",",","np.zeros","(","[","training_steps",",","1","]",")","]","sess.run","(","tf.compat.v1.global_variables_initializer","(",")",")","for","step","in","range","(","training_steps",")",":","sess.run","(","optimizer",")","sample","=","[","sess.run","(","var_1",")",",","sess.run","(","var_2",")","]","samples","[","0","]","[","step",",",":","]","=","sample","[","0","]","samples","[","1","]","[","step",",",":","]","=","sample","[","1","]","samples_","=","np.concatenate","(","samples",",","axis=-1",")","sample_mean","=","np.mean","(","samples_",",","0",")","print","(","'sample","mean","'",",","sample_mean",")","``","`","Args",":","learning_rate",":","Scalar","`","float","`","-like","`","Tensor","`",".","The","base","learning","rate","for","the","optimizer",".","Must","be","tuned","to","the","specific","function","being","minimized",".","preconditioner_decay_rate",":","Scalar","`","float","`","-like","`","Tensor","`",".","The","exponential","decay","rate","of","the","rescaling","of","the","preconditioner","(","RMSprop",")",".","(","This","is","``","alpha","''","in","Li","et","al",".","(","2016",")",")",".","Should","be","smaller","than","but","nearly","`","1","`","to","approximate","sampling","from","the","posterior",".","(","Default",":","`","0.95","`",")","data_size",":","Scalar","`","int","`","-like","`","Tensor","`",".","The","effective","number","of","points","in","the","data","set",".","Assumes","that","the","loss","is","taken","as","the","mean","over","a","minibatch",".","Otherwise","if","the","sum","was","taken",",","divide","this","number","by","the","batch","size",".","If","a","prior","is","included","in","the","loss","function",",","it","should","be","normalized","by","`","data_size","`",".","Default","value",":","`","1","`",".","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","`",")","diagonal_bias",":","Scalar","`","float","`","-like","`","Tensor","`",".","Term","added","to","the","diagonal","of","the","preconditioner","to","prevent","the","preconditioner","from","degenerating",".","(","Default",":","`","1e-8","`",")","name",":","Python","`","str","`","describing","ops","managed","by","this","function",".","(","Default",":","`","``","StochasticGradientLangevinDynamics","''","`",")","parallel_iterations",":","the","number","of","coordinates","for","which","the","gradients","of","the","preconditioning","matrix","can","be","computed","in","parallel",".","Must","be","a","positive","integer",".","Raises",":","InvalidArgumentError",":","If","preconditioner_decay_rate","is","a","`","Tensor","`","not","in","`","(","0,1","]","`",".","NotImplementedError",":","If","eager","execution","is","enabled",".","#","#","#","#","References","[","1","]",":","Chunyuan","Li",",","Changyou","Chen",",","David","Carlson",",","and","Lawrence","Carin",".","Preconditioned","Stochastic","Gradient","Langevin","Dynamics","for","Deep","Neural","Networks",".","In","_Association","for","the","Advancement","of","Artificial","Intelligence_",",","2016.","https",":","\/\/arxiv.org\/abs\/1512.07666"]
docstr: Variable scope of all calls to `tf.get_variable`.
docstr_tok: ["Variable","scope","of","all","calls","to","`","tf.get_variable","`","."]
code: def variable_scope(self):
      """Variable scope of all calls to `tf.get_variable`."""
      return self._variable_scope
  
file_path: probability/tensorflow_probability/python/optimizer/sgld.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
  from tensorflow_probability.python.math.diag_jacobian import diag_jacobian
label_desc: Use image node_id 7 for calling the StochasticGradientLangevinDynamics obj's underlying member method code with example usage: obj.variable_scope() and returns: self