code: # Copyright 2018 The TensorFlow Probability Authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ============================================================================
  """Replica Exchange Monte Carlo Transition Kernel."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import collections
  
  import tensorflow.compat.v1 as tf1
  import tensorflow.compat.v2 as tf
  
  from tensorflow_probability.python.mcmc import kernel as kernel_base
  from tensorflow_probability.python.mcmc.internal import util as mcmc_util
  from tensorflow_probability.python.util.seed_stream import SeedStream
  
  __all__ = [
      'ReplicaExchangeMC',
      'default_exchange_proposed_fn',
  ]
  
  ReplicaExchangeMCKernelResults = collections.namedtuple(
      'ReplicaExchangeMCKernelResults',
      [
          # List of states for each replica.  Each state may itself be a list of
          # state parts.
          'replica_states',
          # List of KernelResults for each replica, post exchange.
          'replica_results',
          # List of state/state-parts with pre-exchange samples from each replica.
          'sampled_replica_states',
          # List of kernel-results with pre-exchange samples from each replica.
          'sampled_replica_results',
      ])
  
  
  def default_exchange_proposed_fn(prob_exchange):
    """Default exchange proposal function, for replica exchange MC.
  
    With probability `prob_exchange` propose combinations of replica for exchange.
    When exchanging, create combinations of adjacent replicas in
    [Replica Exchange Monte Carlo](
    https://en.wikipedia.org/wiki/Parallel_tempering)
  
    ```
    exchange_fn = default_exchange_proposed_fn(prob_exchange=0.5)
    exchange_proposed = exchange_fn(num_replica=3)
  
    exchange_proposed.eval()
    ==> [[0, 1]]  # 1 exchange, 0 <--> 1
  
    exchange_proposed.eval()
    ==> []  # 0 exchanges
    ```
  
    Args:
      prob_exchange: Scalar `Tensor` giving probability that any exchanges will
        be generated.
  
    Returns:
      default_exchange_proposed_fn_: Python callable which take a number of
        replicas (a Python integer), and return combinations of replicas for
        exchange as an [n, 2] integer `Tensor`, `0 <= n <= num_replica // 2`,
        with *unique* values in the set `{0, ..., num_replica}`.
    """
  
    def default_exchange_proposed_fn_(num_replica, seed=None):
      """Default function for `exchange_proposed_fn` of `kernel`."""
      seed_stream = SeedStream(seed, 'default_exchange_proposed_fn')
  
      zero_start = tf.random.uniform([], seed=seed_stream()) > 0.5
      if num_replica % 2 == 0:
  
        def _exchange():
          flat_exchange = tf.range(num_replica)
          if num_replica > 2:
            start = tf.cast(~zero_start, dtype=tf.int32)
            end = num_replica - start
            flat_exchange = flat_exchange[start:end]
          return tf.reshape(flat_exchange, [tf.size(input=flat_exchange) // 2, 2])
      else:
  
        def _exchange():
          start = tf.cast(zero_start, dtype=tf.int32)
          end = num_replica - tf.cast(~zero_start, dtype=tf.int32)
          flat_exchange = tf.range(num_replica)[start:end]
          return tf.reshape(flat_exchange, [tf.size(input=flat_exchange) // 2, 2])
  
      def _null_exchange():
        return tf.reshape(tf.cast([], dtype=tf.int32), shape=[0, 2])
  
      return tf.cond(
          pred=tf.random.uniform([], seed=seed_stream()) < prob_exchange,
          true_fn=_exchange,
          false_fn=_null_exchange)
  
    return default_exchange_proposed_fn_
  
  
  class ReplicaExchangeMC(kernel_base.TransitionKernel):
    """Runs one step of the Replica Exchange Monte Carlo.
  
    [Replica Exchange Monte Carlo](
    https://en.wikipedia.org/wiki/Parallel_tempering) is a Markov chain
    Monte Carlo (MCMC) algorithm that is also known as Parallel Tempering. This
    algorithm performs multiple sampling with different temperatures in parallel,
    and exchanges those samplings according to the Metropolis-Hastings criterion.
  
    The `K` replicas are parameterized in terms of `inverse_temperature`'s,
    `(beta[0], beta[1], ..., beta[K-1])`.  If the target distribution has
    probability density `p(x)`, the `kth` replica has density `p(x)**beta_k`.
  
    Typically `beta[0] = 1.0`, and `1.0 > beta[1] > beta[2] > ... > 0.0`.
  
    * `beta[0] == 1` ==> First replicas samples from the target density, `p`.
    * `beta[k] < 1`, for `k = 1, ..., K-1` ==> Other replicas sample from
      "flattened" versions of `p` (peak is less high, valley less low).  These
      distributions are somewhat closer to a uniform on the support of `p`.
  
    Samples from adjacent replicas `i`, `i + 1` are used as proposals for each
    other in a Metropolis step.  This allows the lower `beta` samples, which
    explore less dense areas of `p`, to occasionally be used to help the
    `beta == 1` chain explore new regions of the support.
  
    Samples from replica 0 are returned, and the others are discarded.
  
    #### Examples
  
    ##### Sampling from the Standard Normal Distribution.
  
    ```python
    import numpy as np
    import tensorflow as tf
    import tensorflow_probability as tfp
    tfd = tfp.distributions
  
    dtype = np.float32
  
    target = tfd.Normal(loc=dtype(0), scale=dtype(1))
  
    def make_kernel_fn(target_log_prob_fn, seed):
      return tfp.mcmc.HamiltonianMonteCarlo(
          target_log_prob_fn=target_log_prob_fn,
          seed=seed, step_size=1.0, num_leapfrog_steps=3)
  
    remc = tfp.mcmc.ReplicaExchangeMC(
        target_log_prob_fn=target.log_prob,
        inverse_temperatures=[1., 0.3, 0.1, 0.03],
        make_kernel_fn=make_kernel_fn,
        seed=42)
  
    samples, _ = tfp.mcmc.sample_chain(
        num_results=1000,
        current_state=dtype(1),
        kernel=remc,
        num_burnin_steps=500,
        parallel_iterations=1)  # For determinism.
  
    sample_mean = tf.reduce_mean(samples, axis=0)
    sample_std = tf.sqrt(
        tf.reduce_mean(tf.squared_difference(samples, sample_mean),
                       axis=0))
    with tf.Session() as sess:
      [sample_mean_, sample_std_] = sess.run([sample_mean, sample_std])
  
    print('Estimated mean: {}'.format(sample_mean_))
    print('Estimated standard deviation: {}'.format(sample_std_))
    ```
  
    ##### Sampling from a 2-D Mixture Normal Distribution.
  
    ```python
    import numpy as np
    import tensorflow as tf
    import tensorflow_probability as tfp
    import matplotlib.pyplot as plt
    tfd = tfp.distributions
  
    dtype = np.float32
  
    target = tfd.MixtureSameFamily(
        mixture_distribution=tfd.Categorical(probs=[0.5, 0.5]),
        components_distribution=tfd.MultivariateNormalDiag(
            loc=[[-1., -1], [1., 1.]],
            scale_identity_multiplier=[0.1, 0.1]))
  
    def make_kernel_fn(target_log_prob_fn, seed):
      return tfp.mcmc.HamiltonianMonteCarlo(
          target_log_prob_fn=target_log_prob_fn,
          seed=seed, step_size=0.3, num_leapfrog_steps=3)
  
    remc = tfp.mcmc.ReplicaExchangeMC(
        target_log_prob_fn=target.log_prob,
        inverse_temperatures=[1., 0.3, 0.1, 0.03, 0.01],
        make_kernel_fn=make_kernel_fn,
        seed=42)
  
    samples, _ = tfp.mcmc.sample_chain(
        num_results=1000,
        # Start near the [1, 1] mode.  Standard HMC would get stuck there.
        current_state=np.ones(2, dtype=dtype),
        kernel=remc,
        num_burnin_steps=500,
        parallel_iterations=1)  # For determinism.
  
    with tf.Session() as sess:
      samples_ = sess.run(samples)
  
    plt.figure(figsize=(8, 8))
    plt.xlim(-2, 2)
    plt.ylim(-2, 2)
    plt.plot(samples_[:, 0], samples_[:, 1], '.')
    plt.show()
    ```
  
    """
  
    def __init__(self,
                 target_log_prob_fn,
                 inverse_temperatures,
                 make_kernel_fn,
                 exchange_proposed_fn=default_exchange_proposed_fn(1.),
                 seed=None,
                 name=None):
      """Instantiates this object.
  
      Args:
        target_log_prob_fn: Python callable which takes an argument like
          `current_state` (or `*current_state` if it's a list) and returns its
          (possibly unnormalized) log-density under the target distribution.
        inverse_temperatures: `1D` `Tensor of inverse temperatures to perform
          samplings with each replica. Must have statically known `shape`.
          `inverse_temperatures[0]` produces the states returned by samplers,
          and is typically == 1.
        make_kernel_fn: Python callable which takes target_log_prob_fn and seed
          args and returns a TransitionKernel instance.
        exchange_proposed_fn: Python callable which take a number of replicas, and
          return combinations of replicas for exchange.
        seed: Python integer to seed the random number generator.
          Default value: `None` (i.e., no seed).
        name: Python `str` name prefixed to Ops created by this function.
          Default value: `None` (i.e., "remc_kernel").
  
      Raises:
        ValueError: `inverse_temperatures` doesn't have statically known 1D shape.
      """
      inverse_temperatures = tf.convert_to_tensor(
          value=inverse_temperatures, name='inverse_temperatures')
  
      # Note these are static checks, and don't need to be embedded in the graph.
      inverse_temperatures.shape.assert_is_fully_defined()
      inverse_temperatures.shape.assert_has_rank(1)
  
      self._seed_stream = SeedStream(seed, salt=name)
      self._seeded_mcmc = seed is not None
      self._parameters = dict(
          target_log_prob_fn=target_log_prob_fn,
          inverse_temperatures=inverse_temperatures,
          num_replica=tf.compat.dimension_value(inverse_temperatures.shape[0]),
          exchange_proposed_fn=exchange_proposed_fn,
          seed=seed,
          name=name)
      self.replica_kernels = []
      for i in range(self.num_replica):
        self.replica_kernels.append(
            make_kernel_fn(
                target_log_prob_fn=_replica_log_prob_fn(inverse_temperatures[i],
                                                        target_log_prob_fn),
                seed=self._seed_stream()))
  
    @property
    def target_log_prob_fn(self):
      return self._parameters['target_log_prob_fn']
  
    @property
    def inverse_temperatures(self):
      return self._parameters['inverse_temperatures']
  
    @property
    def num_replica(self):
      return self._parameters['num_replica']
  
    @property
    def exchange_proposed_fn(self):
      return self._parameters['exchange_proposed_fn']
  
    @property
    def seed(self):
      return self._parameters['seed']
  
    @property
    def name(self):
      return self._parameters['name']
  
    @property
    def parameters(self):
      """Return `dict` of ``__init__`` arguments and their values."""
      return self._parameters
  
    @property
    def is_calibrated(self):
      return True
  
    def one_step(self, current_state, previous_kernel_results):
      """Takes one step of the TransitionKernel.
  
      Args:
        current_state: `Tensor` or Python `list` of `Tensor`s representing the
          current state(s) of the Markov chain(s).
        previous_kernel_results: A (possibly nested) `tuple`, `namedtuple` or
          `list` of `Tensor`s representing internal calculations made within the
          previous call to this function (or as returned by `bootstrap_results`).
  
      Returns:
        next_state: `Tensor` or Python `list` of `Tensor`s representing the
          next state(s) of the Markov chain(s).
        kernel_results: A (possibly nested) `tuple`, `namedtuple` or `list` of
          `Tensor`s representing internal calculations made within this function.
          This inculdes replica states.
      """
      # Key difficulty:  The type of exchanges differs from one call to the
      # next...even the number of exchanges can differ.
      # As a result, exchanges must happen dynamically, in while loops.
      with tf1.name_scope(
          name=mcmc_util.make_name(self.name, 'remc', 'one_step'),
          values=[current_state, previous_kernel_results]):
  
        # Each replica does `one_step` to get pre-exchange states/KernelResults.
        sampled_replica_states, sampled_replica_results = zip(*[
            rk.one_step(previous_kernel_results.replica_states[i],
                        previous_kernel_results.replica_results[i])
            for i, rk in enumerate(self.replica_kernels)
        ])
        sampled_replica_states = list(sampled_replica_states)
        sampled_replica_results = list(sampled_replica_results)
  
        states_are_lists = mcmc_util.is_list_like(sampled_replica_states[0])
  
        if not states_are_lists:
          sampled_replica_states = [[s] for s in sampled_replica_states]
        num_state_parts = len(sampled_replica_states[0])
  
        dtype = sampled_replica_states[0][0].dtype
  
        # Must put states into TensorArrays.  Why?  We will read/write states
        # dynamically with Tensor index `i`, and you cannot do this with lists.
        # old_states[k][i] is Tensor of (old) state part k, for replica i.
        # The `k` will be known statically, and `i` is a Tensor.
        old_states = [
            tf.TensorArray(
                dtype,
                size=self.num_replica,
                dynamic_size=False,
                clear_after_read=False,
                tensor_array_name='old_states',
                # State part k has same shape, regardless of replica.  So use 0.
                element_shape=sampled_replica_states[0][k].shape)
            for k in range(num_state_parts)
        ]
        for k in range(num_state_parts):
          for i in range(self.num_replica):
            old_states[k] = old_states[k].write(i, sampled_replica_states[i][k])
  
        exchange_proposed = self.exchange_proposed_fn(
            self.num_replica, seed=self._seed_stream())
        exchange_proposed_n = tf.shape(input=exchange_proposed)[0]
  
        exchanged_states = self._get_exchanged_states(
            old_states, exchange_proposed, exchange_proposed_n,
            sampled_replica_states, sampled_replica_results)
  
        no_exchange_proposed, _ = tf1.setdiff1d(
            tf.range(self.num_replica), tf.reshape(exchange_proposed, [-1]))
  
        exchanged_states = self._insert_old_states_where_no_exchange_was_proposed(
            no_exchange_proposed, old_states, exchanged_states)
  
        next_replica_states = []
        for i in range(self.num_replica):
          next_replica_states_i = []
          for k in range(num_state_parts):
            next_replica_states_i.append(exchanged_states[k].read(i))
          next_replica_states.append(next_replica_states_i)
  
        if not states_are_lists:
          next_replica_states = [s[0] for s in next_replica_states]
          sampled_replica_states = [s[0] for s in sampled_replica_states]
  
        # Now that states are/aren't exchanged, bootstrap next kernel_results.
        # The viewpoint is that after each exchange, we are starting anew.
        next_replica_results = [
            rk.bootstrap_results(state)
            for rk, state in zip(self.replica_kernels, next_replica_states)
        ]
  
        next_state = next_replica_states[0]  # Replica 0 is the returned state(s).
  
        kernel_results = ReplicaExchangeMCKernelResults(
            replica_states=next_replica_states,
            replica_results=next_replica_results,
            sampled_replica_states=sampled_replica_states,
            sampled_replica_results=sampled_replica_results,
        )
  
        return next_state, kernel_results
  
    def _get_exchanged_states(self, old_states, exchange_proposed,
                              exchange_proposed_n, sampled_replica_states,
                              sampled_replica_results):
      """Get list of TensorArrays holding exchanged states, and zeros."""
      with tf1.name_scope('get_exchanged_states'):
  
        target_log_probs = []
        for replica in range(self.num_replica):
          replica_log_prob = _get_field(sampled_replica_results[replica],
                                        'target_log_prob')
          inverse_temp = self.inverse_temperatures[replica]
          target_log_probs.append(replica_log_prob / inverse_temp)
        target_log_probs = tf.stack(target_log_probs, axis=0)
  
        dtype = target_log_probs.dtype
        num_state_parts = len(sampled_replica_states[0])
        # exchanged_states[k][i] is Tensor of (new) state part k, for replica i.
        # The `k` will be known statically, and `i` is a Tensor.
        # We will insert values into indices `i` for every replica with a proposed
        # exchange.
        exchanged_states = [
            tf.TensorArray(
                dtype,
                size=self.num_replica,
                dynamic_size=False,
                tensor_array_name='exchanged_states',
                # State part k has same shape, regardless of replica.  So use 0.
                element_shape=sampled_replica_states[0][k].shape)
            for k in range(num_state_parts)
        ]
  
        # Draw random variables here, to avoid sampling in the loop (and losing
        # reproducibility).  This may mean we sample too many, but we will always
        # have enough.
        sample_shape = tf.concat(
            ([self.num_replica // 2], tf.shape(input=target_log_probs)[1:]),
            axis=0)
        log_uniforms = tf.math.log(
            tf.random.uniform(
                shape=sample_shape, dtype=dtype, seed=self._seed_stream()))
  
        def _swap(is_exchange_accepted, x, y):
          """Swap batches of x, y where accepted."""
          with tf1.name_scope('swap_where_exchange_accepted'):
            new_x = mcmc_util.choose(is_exchange_accepted, y, x)
            new_y = mcmc_util.choose(is_exchange_accepted, x, y)
          return new_x, new_y
  
        def cond(i, unused_exchanged_states):
          return i < exchange_proposed_n
  
        def body(i, exchanged_states):
          """Body of while loop for exchanging states."""
          # Propose exchange between replicas indexed by m and n.
          m, n = tf.unstack(exchange_proposed[i])
  
          # Construct log_accept_ratio:  -temp_diff * target_log_prob_diff.
          # Note target_log_prob_diff = -EnergyDiff (common definition is in terms
          # of energy).
          temp_diff = self.inverse_temperatures[m] - self.inverse_temperatures[n]
          # Difference of target log probs may be +- Inf or NaN.  We want the
          # product of this with the temperature difference to have "alt value" of
          # -Inf.
          log_accept_ratio = mcmc_util.safe_sum(
              [-temp_diff * target_log_probs[m], temp_diff * target_log_probs[n]])
  
          is_exchange_accepted = log_uniforms[i] < log_accept_ratio
  
          for k in range(num_state_parts):
            new_m, new_n = _swap(is_exchange_accepted, old_states[k].read(m),
                                 old_states[k].read(n))
            exchanged_states[k] = exchanged_states[k].write(m, new_m)
            exchanged_states[k] = exchanged_states[k].write(n, new_n)
  
          return i + 1, exchanged_states
  
        # At this point, exchanged_states[k] is a length num_replicas TensorArray.
        return tf.while_loop(
            cond=cond, body=body, loop_vars=[tf.constant(0),
                                             exchanged_states])[1]  # Remove `i`
  
    def _insert_old_states_where_no_exchange_was_proposed(
        self, no_exchange_proposed, old_states, exchanged_states):
      with tf1.name_scope(
          'insert_old_states_where_no_exchange_was_proposed'):
  
        def cond(j, unused_exchanged_states):
          return j < tf.size(input=no_exchange_proposed)
  
        def body(j, exchanged_states):
          replica = no_exchange_proposed[j]
          for k in range(len(old_states)):  # k indexes state part
            exchanged_states[k] = exchanged_states[k].write(
                replica, old_states[k].read(replica))
          return j + 1, exchanged_states
  
        return tf.while_loop(
            cond=cond, body=body, loop_vars=[tf.constant(0),
                                             exchanged_states])[1]  # Remove `j`
  
    def bootstrap_results(self, init_state):
      """Returns an object with the same type as returned by `one_step`.
  
      Args:
        init_state: `Tensor` or Python `list` of `Tensor`s representing the
          initial state(s) of the Markov chain(s).
  
      Returns:
        kernel_results: A (possibly nested) `tuple`, `namedtuple` or `list` of
          `Tensor`s representing internal calculations made within this function.
          This inculdes replica states.
      """
      with tf1.name_scope(
          name=mcmc_util.make_name(self.name, 'remc', 'bootstrap_results'),
          values=[init_state]):
        replica_results = [
            self.replica_kernels[i].bootstrap_results(init_state)
            for i in range(self.num_replica)
        ]
  
        init_state_parts = (
            list(init_state)
            if mcmc_util.is_list_like(init_state) else [init_state])
  
        # Convert all states parts to tensor...
        replica_states = [[
            tf.convert_to_tensor(value=s) for s in init_state_parts
        ] for i in range(self.num_replica)]
  
        if not mcmc_util.is_list_like(init_state):
          replica_states = [s[0] for s in replica_states]
  
        return ReplicaExchangeMCKernelResults(
            replica_states=replica_states,
            replica_results=replica_results,
            sampled_replica_states=replica_states,
            sampled_replica_results=replica_results,
        )
  
  
  def _replica_log_prob_fn(inverse_temperature, target_log_prob_fn):
    """Return a log probability function made considering temperature."""
  
    def _replica_log_prob_fn_(*x):
      return inverse_temperature * target_log_prob_fn(*x)
  
    return _replica_log_prob_fn_
  
  
  # TODO(b/111801087) Use a more standardized API when available.
  def _get_field(kernel_results, field_name):
    """field_name from kernel_results or kernel_results.accepted_results."""
    if hasattr(kernel_results, field_name):
      return getattr(kernel_results, field_name)
    if hasattr(kernel_results, 'accepted_results'):
      return getattr(kernel_results.accepted_results, field_name)
    raise TypeError('Cannot extract %s from %s' % (field_name, kernel_results))
   # Copyright 2018 The TensorFlow Probability Authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ============================================================================
  """Functional test for GradientDescent."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  # Dependency imports
  
  import numpy as np
  import tensorflow.compat.v1 as tf1
  import tensorflow.compat.v2 as tf
  import tensorflow_probability as tfp
  from tensorflow_probability.python import distributions as tfd
  from tensorflow_probability.python.internal import test_case
  from tensorflow_probability.python.math import diag_jacobian
  
  from tensorflow.python.framework import test_util  # pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
  
  
  @test_util.run_all_in_graph_and_eager_modes
  class StochasticGradientLangevinDynamicsOptimizerTest(test_case.TestCase):
  
    def testBasic(self):
      if 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., 4.], dtype=dtype)
          grads0 = tf.constant([0.1, 0.1], dtype=dtype)
          grads1 = tf.constant([0.01, 0.01], dtype=dtype)
          decay_rate = 0.53
          sgd_optimizer = tfp.optimizer.StochasticGradientLangevinDynamics(
              3., preconditioner_decay_rate=decay_rate)
          sgd_op = sgd_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., 4.], self.evaluate(var1))
          # Run 1 step of sgd
          self.evaluate(sgd_op)
          # Validate updated params
          grads_scaled = (0.5 * 0.1 /
                          np.sqrt(decay_rate + (1. - decay_rate) * 0.1**2 + 1e-8))
          # Note that `tfp.math.diag_jacobian(xs=var, ys=grad)` returns zero
          # tensor
          self.assertAllCloseAccordingToType(
              [1.1 - 3. * grads_scaled, 2.1 - 3. * grads_scaled],
              self.evaluate(var0))
          grads_scaled = (0.5 * 0.01 / np.sqrt(
              decay_rate + (1. - decay_rate) * 0.01**2 + 1e-8))
          self.assertAllCloseAccordingToType(
              [3. - 3. * grads_scaled, 4. - 3. * grads_scaled],
              self.evaluate(var1))
          self.assertAllCloseAccordingToType(
              1, self.evaluate(sgd_optimizer.iterations))
  
    def testBasicMultiInstance(self):
      if 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., 4.], dtype=dtype)
          grads0 = tf.constant([0.1, 0.1], dtype=dtype)
          grads1 = tf.constant([0.01, 0.01], dtype=dtype)
          vara = tf.Variable([1.1, 2.1], dtype=dtype)
          varb = tf.Variable([3., 4.], dtype=dtype)
          gradsa = tf.constant([0.1, 0.1], dtype=dtype)
          gradsb = tf.constant([0.01, 0.01], dtype=dtype)
          decay_rate = 0.5
          sgd_optimizer = tfp.optimizer.StochasticGradientLangevinDynamics(
              3., preconditioner_decay_rate=decay_rate)
          sgd_op = sgd_optimizer.apply_gradients(
              zip([grads0, grads1], [var0, var1]))
          sgd_optimizer2 = tfp.optimizer.StochasticGradientLangevinDynamics(
              3., preconditioner_decay_rate=decay_rate)
          sgd_op2 = sgd_optimizer2.apply_gradients(
              zip([gradsa, gradsb], [vara, varb]))
          self.evaluate(tf1.global_variables_initializer())
          # Fetch params to validate initial values
          self.assertAllCloseAccordingToType([1.1, 2.1], self.evaluate(var0))
          self.assertAllCloseAccordingToType([3., 4.], self.evaluate(var1))
          self.assertAllCloseAccordingToType([1.1, 2.1], self.evaluate(vara))
          self.assertAllCloseAccordingToType([3., 4.], self.evaluate(varb))
  
          # Run 1 step of sgd
          self.evaluate(sgd_op)
          self.evaluate(sgd_op2)
  
          # Validate updated params
          grads_scaled = (0.5 * 0.1 /
                          np.sqrt(decay_rate + (1. - decay_rate) * 0.1**2 + 1e-8))
          self.assertAllCloseAccordingToType(
              [1.1 - 3. * grads_scaled, 2.1 - 3. * grads_scaled],
              self.evaluate(var0))
          self.assertAllCloseAccordingToType(
              [1.1 - 3. * grads_scaled, 2.1 - 3. * grads_scaled],
              self.evaluate(vara))
  
          grads_scaled = (0.5 * 0.01 / np.sqrt(
              decay_rate + (1 - decay_rate) * 0.01**2 + 1e-8))
          self.assertAllCloseAccordingToType(
              [3. - 3. * grads_scaled, 4. - 3. * grads_scaled],
              self.evaluate(var1))
          self.assertAllCloseAccordingToType(
              [3. - 3. * grads_scaled, 4. - 3. * grads_scaled],
              self.evaluate(varb))
          self.assertAllCloseAccordingToType(
              1, self.evaluate(sgd_optimizer.iterations))
          self.assertAllCloseAccordingToType(
              1, self.evaluate(sgd_optimizer2.iterations))
  
    def testTensorLearningRate(self):
      if 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., 4.], dtype=dtype)
          grads0 = tf.constant([0.1, 0.1], dtype=dtype)
          grads1 = tf.constant([0.01, 0.01], dtype=dtype)
          lrate = tf.constant(3.0)
          decay_rate = 0.5
          sgd_op = tfp.optimizer.StochasticGradientLangevinDynamics(
              lrate, preconditioner_decay_rate=tf.constant(
                  decay_rate)).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., 4.], self.evaluate(var1))
          # Run 1 step of sgd
          self.evaluate(sgd_op)
          # Validate updated params
          grads_scaled = (0.5 * 0.1 /
                          np.sqrt(decay_rate + (1. - decay_rate) * 0.1**2 + 1e-8))
          # Note that `tfp.math.diag_jacobian(xs=var, ys=grad)` returns zero
          # tensor
          self.assertAllCloseAccordingToType(
              [1.1 - 3. * grads_scaled, 2.1 - 3. * grads_scaled],
              self.evaluate(var0))
          grads_scaled = (0.5 * 0.01 / np.sqrt(
              decay_rate + (1. - decay_rate) * 0.01**2 + 1e-8))
          self.assertAllCloseAccordingToType(
              [3. - 3. * grads_scaled, 4. - 3. * grads_scaled],
              self.evaluate(var1))
  
    @test_util.run_deprecated_v1
    def testGradWrtRef(self):
      if tf.executing_eagerly():
        return
  
      for dtype in [tf.half, tf.float32, tf.float64]:
        with self.cached_session():
          opt = tfp.optimizer.StochasticGradientLangevinDynamics(3.0)
          values = [1., 3.]
          vars_ = [tf.Variable([v], dtype=dtype) for v in values]
          loss = lambda: vars_[0] + vars_[1]  # pylint: disable=cell-var-from-loop
          grads_and_vars = opt._compute_gradients(loss, vars_)
          self.evaluate(tf1.global_variables_initializer())
          for grad, _ in grads_and_vars:
            self.assertAllCloseAccordingToType([1.], self.evaluate(grad))
  
    def testBurnin(self):
      if tf.executing_eagerly():
        return
  
      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 = tfp.optimizer.StochasticGradientLangevinDynamics(
              3.,
              preconditioner_decay_rate=decay_rate,
              burnin=tf.constant(10, dtype=burnin_dtype))
          sgd_op = sgd_optimizer.apply_gradients([(grads0, var0)])
  
          self.evaluate(tf1.global_variables_initializer())
          # Validate that iterations is initialized to 0.
          self.assertAllCloseAccordingToType(
              0, self.evaluate(sgd_optimizer.iterations))
          # Run 1 step of sgd
          self.evaluate(sgd_op)
          # Validate that iterations is incremented.
          self.assertAllCloseAccordingToType(
              1, self.evaluate(sgd_optimizer.iterations))
  
    def testWithGlobalStep(self):
      if tf.executing_eagerly():
        return
  
      for dtype in [tf.float32, tf.float64]:
        with self.cached_session():
          step = tf.Variable(0, dtype=tf.int64)
  
          var0 = tf.Variable([1.1, 2.1], dtype=dtype)
          var1 = tf.Variable([3., 4.], dtype=dtype)
          grads0 = tf.constant([0.1, 0.1], dtype=dtype)
          grads1 = tf.constant([0.01, 0.01], dtype=dtype)
          decay_rate = 0.1
  
          sgd_opt = tfp.optimizer.StochasticGradientLangevinDynamics(
              3., preconditioner_decay_rate=decay_rate)
          sgd_opt.iterations = step
          sgd_op = sgd_opt.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., 4.], self.evaluate(var1))
          # Run 1 step of sgd
          self.evaluate(sgd_op)
  
          # Validate updated params and step
          grads_scaled = (0.5 * 0.1 /
                          np.sqrt(decay_rate + (1. - decay_rate) * 0.1**2 + 1e-8))
          # Note that `tfp.math.diag_jacobian(xs=var, ys=grad)` returns zero
          # tensor
          self.assertAllCloseAccordingToType(
              [1.1 - 3. * grads_scaled, 2.1 - 3. * grads_scaled],
              self.evaluate(var0))
          grads_scaled = (0.5 * 0.01 / np.sqrt(
              decay_rate + (1. - decay_rate) * 0.01**2 + 1e-8))
          self.assertAllCloseAccordingToType(
              [3. - 3. * grads_scaled, 4. - 3. * grads_scaled],
              self.evaluate(var1))
          self.assertAllCloseAccordingToType(1, self.evaluate(step))
  
    def testSparseBasic(self):
      if 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.], [4.]], dtype=dtype)
          grads0 = tf.IndexedSlices(
              tf.constant([0.1], shape=[1, 1], dtype=dtype),
              tf.constant([0]), tf.constant([2, 1]))
          grads1 = tf.IndexedSlices(
              tf.constant([0.01], shape=[1, 1], dtype=dtype),
              tf.constant([1]), tf.constant([2, 1]))
          decay_rate = 0.9
          sgd_op = tfp.optimizer.StochasticGradientLangevinDynamics(
              3., preconditioner_decay_rate=decay_rate).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.], [4.]], self.evaluate(var1))
          # Run 1 step of sgd
          self.evaluate(sgd_op)
          # Validate updated params
          grads_scaled = (0.5 * 0.1 /
                          np.sqrt(decay_rate + (1. - decay_rate) * 0.1**2 + 1e-8))
          # Note that `tfp.math.diag_jacobian(xs=var, ys=grad)` returns zero
          # tensor
          self.assertAllCloseAccordingToType([[1.1 - 3. * grads_scaled], [2.1]],
                                             self.evaluate(var0))
          grads_scaled = (0.5 * 0.01 / np.sqrt(
              decay_rate + (1. - decay_rate) * 0.01**2 + 1e-8))
          self.assertAllCloseAccordingToType(
              [[3. - 3. * 0], [4. - 3. * grads_scaled]], self.evaluate(var1))
  
    def testPreconditionerComputedCorrectly(self):
      """Test that SGLD step is computed correctly for a 3D Gaussian energy."""
      if tf.executing_eagerly():
        return
  
      with self.cached_session():
        dtype = np.float32
        # Target function is the energy function of normal distribution
        true_mean = dtype([0, 0, 0])
        true_cov = dtype([[1, 0.25, 0.25], [0.25, 1, 0.25], [0.25, 0.25, 1]])
        # Target distribution is defined through the Cholesky decomposition
        chol = tf.linalg.cholesky(true_cov)
        target = tfd.MultivariateNormalTriL(loc=true_mean, scale_tril=chol)
        var_1 = tf.Variable(name='var_1', initial_value=[1., 1.])
        var_2 = tf.Variable(name='var_2', initial_value=[1.])
  
        var = [var_1, var_2]
  
        # Set up the learning rate and the optimizer
        learning_rate = .5
        optimizer_kernel = tfp.optimizer.StochasticGradientLangevinDynamics(
            learning_rate=learning_rate, burnin=1)
  
        # Target function
        def target_fn(x, y):
          # Stack the input tensors together
          z = tf.concat([x, y], axis=-1) - true_mean
          return -target.log_prob(z)
  
        grads = tf.gradients(ys=target_fn(*var), xs=var)
  
        # Update value of `var` with one iteration of the SGLD (without the
        # normal perturbation, since `burnin > 0`)
        step = optimizer_kernel.apply_gradients(zip(grads, var))
  
        # True theoretical value of `var` after one iteration
        decay_tensor = tf.cast(optimizer_kernel._decay_tensor, var[0].dtype)
        diagonal_bias = tf.cast(optimizer_kernel._diagonal_bias, var[0].dtype)
        learning_rate = tf.cast(optimizer_kernel._learning_rate, var[0].dtype)
        velocity = [(decay_tensor * tf.ones_like(v)
                     + (1 - decay_tensor) * tf.square(g))
                    for v, g in zip(var, grads)]
        preconditioner = [tf.math.rsqrt(vel + diagonal_bias) for vel in velocity]
        # Compute second order gradients
        _, grad_grads = diag_jacobian(
            xs=var,
            ys=grads)
        # Compute gradient of the preconditioner (compute the gradient manually)
        preconditioner_grads = [-(g * g_g * (1. - decay_tensor) * p**3.)
                                for g, g_g, p in zip(grads, grad_grads,
                                                     preconditioner)]
  
        # True theoretical value of `var` after one iteration
        var_true = [v - learning_rate * 0.5 * (p * g - p_g)
                    for v, p, g, p_g in zip(var, preconditioner, grads,
                                            preconditioner_grads)]
        self.evaluate(tf1.global_variables_initializer())
        var_true_ = self.evaluate(var_true)
        self.evaluate(step)
        var_ = self.evaluate(var)  # new `var` after one SGLD step
        self.assertAllClose(var_true_,
                            var_, atol=0.001, rtol=0.001)
  
    def testDiffusionBehavesCorrectly(self):
      """Test that for the SGLD finds minimum of the 3D Gaussian energy."""
      if tf.executing_eagerly():
        return
  
      with self.cached_session():
        # Set up random seed for the optimizer
        tf1.set_random_seed(42)
        dtype = np.float32
        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.])
  
        # Loss function
        def loss_fn():
          var = tf.concat([var_1, var_2], axis=-1)
          loss_part = tf.linalg.cholesky_solve(chol, tf.expand_dims(var, -1))
          return tf.linalg.matvec(loss_part, var, transpose_a=True)
  
        # Set up the learning rate with a polynomial decay
        global_step = tf1.train.get_or_create_global_step()
        starter_learning_rate = .3
        end_learning_rate = 1e-4
        decay_steps = 1e4
        learning_rate = tf1.train.polynomial_decay(
            starter_learning_rate,
            global_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 = global_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])]
        self.evaluate(tf1.global_variables_initializer())
        for step in range(training_steps):
          self.evaluate(optimizer)
          sample = [self.evaluate(var_1), self.evaluate(var_2)]
          samples[0][step, :] = sample[0]
          samples[1][step, :] = sample[1]
  
      samples_ = np.concatenate(samples, axis=-1)
      sample_mean = np.mean(samples_, 0)
      self.assertAllClose(sample_mean, true_mean, atol=0.1, rtol=0.1)
  
  if __name__ == '__main__':
    tf.test.main()
  
apis: tensorflow.compat.v2.stack tensorflow.compat.v2.shape tensorflow.compat.v2.size tensorflow.compat.v2.unstack tensorflow.compat.v2.convert_to_tensor tensorflow.compat.v2.reshape tensorflow.compat.v2.range tensorflow.compat.v2.TensorArray tensorflow.compat.v2.compat.dimension_value tensorflow.compat.v2.constant tensorflow.compat.v2.cast tensorflow.compat.v1.name_scope
apis: tensorflow.compat.v2.concat tensorflow.compat.v2.square tensorflow.compat.v2.ones_like tensorflow.compat.v2.constant tensorflow.compat.v2.linalg.cholesky tensorflow.compat.v2.expand_dims tensorflow.compat.v2.cast tensorflow.compat.v2.Variable numpy.mean tensorflow.compat.v1.train.polynomial_decay numpy.zeros tensorflow.compat.v2.math.rsqrt tensorflow.compat.v1.global_variables_initializer tensorflow.compat.v2.executing_eagerly tensorflow.compat.v2.linalg.matvec tensorflow.compat.v2.test.main tensorflow.compat.v1.set_random_seed numpy.sqrt numpy.concatenate tensorflow.compat.v1.train.get_or_create_global_step