class_docstr: 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 takes multiple samples (from tempered distributions) in parallel,
  then swaps these samples according to the Metropolis-Hastings criterion.
  See also the review paper [1].
  
  The `K` replicas are parameterized in terms of `inverse_temperature`'s,
  `(beta[0], beta[1], ..., beta[K-1])`. If the user provides
  `target_log_prob_fn`, then the `kth` replica samples from density `p_k(x)`,
  with `log(p_k(x)) = beta_k * target_log_prob(x)`.
  In this case, geometrically decaying `beta` often works well.  That is, with
  `R < 1`, we recommend trying `beta[k] = R^k` so that
  `1.0 = beta[0] > beta[1] > ... > 0`.  See [2].
  
  The user can also provide two functions, `tempered_log_prob_fn` and
  `untempered_log_prob_fn`. In this case, the `kth` replica samples from
  density `p_k(x)` with
  `log(p_k(x)) = beta_k * tempered_log_prob_fn(x) + untempered_log_prob_fn(x)`.
  In this case, `beta` may be zero, and one often sets `beta[-1]` to zero.
  This means the last replica samples using `untempered_log_prob_fn`.
  In the Bayesian setup, `untempered_log_prob_fn` will often be the log prior,
  and `tempered_log_prob_fn` the likelihood.
  
  In all cases,
  
  * `beta[0] == 1` ==> First replica samples from the target density.
  * `beta[k] < 1`, for `k = 1, ..., K-1` ==> Other replicas sample from
    "tempered" versions of target (peak is less high, valley less low).  These
    distributions should allow easier exploration of separated modes.
  
  By default, 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 eventually swap state with the
  `beta == 1` chain, allowing it to explore these new regions.
  
  Samples from replica 0 are returned, and the others are discarded, unless
  `state_includes_replicas`.
  
  #### 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))
  
  # Geometric decay is a good rule of thumb.
  inverse_temperatures = 0.5**tf.range(4, dtype=dtype)
  
  # If everything was Normal, step_size should be ~ sqrt(temperature).
  step_size = 1.5 / tf.sqrt(inverse_temperatures)
  
  def make_kernel_fn(target_log_prob_fn):
    return tfp.mcmc.HamiltonianMonteCarlo(
        target_log_prob_fn=target_log_prob_fn,
        step_size=step_size, num_leapfrog_steps=3)
  
  remc = tfp.mcmc.ReplicaExchangeMC(
      target_log_prob_fn=target.log_prob,
      inverse_temperatures=inverse_temperatures,
      make_kernel_fn=make_kernel_fn)
  
  def trace_swaps(unused_state, results):
    return (results.is_swap_proposed_adjacent,
            results.is_swap_accepted_adjacent)
  
  samples, (is_swap_proposed_adjacent, is_swap_accepted_adjacent) = (
      tfp.mcmc.sample_chain(
          num_results=1000,
          current_state=1.0,
          kernel=remc,
          num_burnin_steps=500,
          trace_fn=trace_swaps)
  )
  
  # conditional_swap_prob[k] = P[ExchangeAccepted | ExchangeProposed],
  # for the swap between replicas k and k+1.
  conditional_swap_prob = (
      tf.reduce_sum(tf.cast(is_swap_accepted_adjacent, tf.float32), axis=0)
      /
      tf.reduce_sum(tf.cast(is_swap_proposed_adjacent, tf.float32), axis=0))
  ```
  
  ##### 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_diag=0.1*tf.ones([2, 2])))
  
  inverse_temperatures = 0.2**tf.range(4, dtype=dtype)
  
  # step_size must broadcast with all batch and event dimensions of target.
  # Here, this means it must broadcast with:
  #  [len(inverse_temperatures)] + target.event_shape
  step_size = 0.075 / tf.reshape(tf.sqrt(inverse_temperatures), shape=(4, 1))
  
  def make_kernel_fn(target_log_prob_fn):
    return tfp.mcmc.HamiltonianMonteCarlo(
        target_log_prob_fn=target_log_prob_fn,
        step_size=step_size, num_leapfrog_steps=3)
  
  remc = tfp.mcmc.ReplicaExchangeMC(
      target_log_prob_fn=target.log_prob,
      inverse_temperatures=inverse_temperatures,
      make_kernel_fn=make_kernel_fn)
  
  samples = tfp.mcmc.sample_chain(
      num_results=1000,
      # Start near the [1, 1] mode. Standard HMC would get stuck there.
      current_state=tf.ones(2, dtype=dtype),
      kernel=remc,
      trace_fn=None,
      num_burnin_steps=500)
  
  plt.figure(figsize=(8, 8))
  plt.xlim(-2, 2)
  plt.ylim(-2, 2)
  plt.plot(samples[:, 0], samples[:, 1], '.')
  plt.show()
  ```
  
  #### References
  
  [1]: David J. Earl, Michael W. Deem
       Parallel Tempering: Theory, Applications, and New Perspectives
       https://arxiv.org/abs/physics/0508111
  [2]: David A. Kofke
       On the acceptance probability of replica-exchange Monte Carlo trials.
       J. of Chem. Phys. Vol. 117 No. 5.
class_docstr_tok: ["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","takes","multiple","samples","(","from","tempered","distributions",")","in","parallel",",","then","swaps","these","samples","according","to","the","Metropolis-Hastings","criterion",".","See","also","the","review","paper","[","1","]",".","The","`","K","`","replicas","are","parameterized","in","terms","of","`","inverse_temperature","`","'s",",","`","(","beta","[","0","]",",","beta","[","1","]",",","...",",","beta","[","K-1","]",")","`",".","If","the","user","provides","`","target_log_prob_fn","`",",","then","the","`","kth","`","replica","samples","from","density","`","p_k","(","x",")","`",",","with","`","log","(","p_k","(","x",")",")","=","beta_k","*","target_log_prob","(","x",")","`",".","In","this","case",",","geometrically","decaying","`","beta","`","often","works","well",".","That","is",",","with","`","R","<","1","`",",","we","recommend","trying","`","beta","[","k","]","=","R^k","`","so","that","`","1.0","=","beta","[","0","]",">","beta","[","1","]",">","...",">","0","`",".","See","[","2","]",".","The","user","can","also","provide","two","functions",",","`","tempered_log_prob_fn","`","and","`","untempered_log_prob_fn","`",".","In","this","case",",","the","`","kth","`","replica","samples","from","density","`","p_k","(","x",")","`","with","`","log","(","p_k","(","x",")",")","=","beta_k","*","tempered_log_prob_fn","(","x",")","+","untempered_log_prob_fn","(","x",")","`",".","In","this","case",",","`","beta","`","may","be","zero",",","and","one","often","sets","`","beta","[","-1","]","`","to","zero",".","This","means","the","last","replica","samples","using","`","untempered_log_prob_fn","`",".","In","the","Bayesian","setup",",","`","untempered_log_prob_fn","`","will","often","be","the","log","prior",",","and","`","tempered_log_prob_fn","`","the","likelihood",".","In","all","cases",",","*","`","beta","[","0","]","==","1","`","==",">","First","replica","samples","from","the","target","density",".","*","`","beta","[","k","]","<","1","`",",","for","`","k","=","1",",","...",",","K-1","`","==",">","Other","replicas","sample","from","``","tempered","''","versions","of","target","(","peak","is","less","high",",","valley","less","low",")",".","These","distributions","should","allow","easier","exploration","of","separated","modes",".","By","default",",","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","eventually","swap","state","with","the","`","beta","==","1","`","chain",",","allowing","it","to","explore","these","new","regions",".","Samples","from","replica","0","are","returned",",","and","the","others","are","discarded",",","unless","`","state_includes_replicas","`",".","#","#","#","#","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",")",")","#","Geometric","decay","is","a","good","rule","of","thumb",".","inverse_temperatures","=","0.5","*","*","tf.range","(","4",",","dtype=dtype",")","#","If","everything","was","Normal",",","step_size","should","be","~","sqrt","(","temperature",")",".","step_size","=","1.5","\/","tf.sqrt","(","inverse_temperatures",")","def","make_kernel_fn","(","target_log_prob_fn",")",":","return","tfp.mcmc.HamiltonianMonteCarlo","(","target_log_prob_fn=target_log_prob_fn",",","step_size=step_size",",","num_leapfrog_steps=3",")","remc","=","tfp.mcmc.ReplicaExchangeMC","(","target_log_prob_fn=target.log_prob",",","inverse_temperatures=inverse_temperatures",",","make_kernel_fn=make_kernel_fn",")","def","trace_swaps","(","unused_state",",","results",")",":","return","(","results.is_swap_proposed_adjacent",",","results.is_swap_accepted_adjacent",")","samples",",","(","is_swap_proposed_adjacent",",","is_swap_accepted_adjacent",")","=","(","tfp.mcmc.sample_chain","(","num_results=1000",",","current_state=1.0",",","kernel=remc",",","num_burnin_steps=500",",","trace_fn=trace_swaps",")",")","#","conditional_swap_prob","[","k","]","=","P","[","ExchangeAccepted","|","ExchangeProposed","]",",","#","for","the","swap","between","replicas","k","and","k+1",".","conditional_swap_prob","=","(","tf.reduce_sum","(","tf.cast","(","is_swap_accepted_adjacent",",","tf.float32",")",",","axis=0",")","\/","tf.reduce_sum","(","tf.cast","(","is_swap_proposed_adjacent",",","tf.float32",")",",","axis=0",")",")","``","`","#","#","#","#","#","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_diag=0.1","*","tf.ones","(","[","2",",","2","]",")",")",")","inverse_temperatures","=","0.2","*","*","tf.range","(","4",",","dtype=dtype",")","#","step_size","must","broadcast","with","all","batch","and","event","dimensions","of","target",".","#","Here",",","this","means","it","must","broadcast","with",":","#","[","len","(","inverse_temperatures",")","]","+","target.event_shape","step_size","=","0.075","\/","tf.reshape","(","tf.sqrt","(","inverse_temperatures",")",",","shape=","(","4",",","1",")",")","def","make_kernel_fn","(","target_log_prob_fn",")",":","return","tfp.mcmc.HamiltonianMonteCarlo","(","target_log_prob_fn=target_log_prob_fn",",","step_size=step_size",",","num_leapfrog_steps=3",")","remc","=","tfp.mcmc.ReplicaExchangeMC","(","target_log_prob_fn=target.log_prob",",","inverse_temperatures=inverse_temperatures",",","make_kernel_fn=make_kernel_fn",")","samples","=","tfp.mcmc.sample_chain","(","num_results=1000",",","#","Start","near","the","[","1",",","1","]","mode",".","Standard","HMC","would","get","stuck","there",".","current_state=tf.ones","(","2",",","dtype=dtype",")",",","kernel=remc",",","trace_fn=None",",","num_burnin_steps=500",")","plt.figure","(","figsize=","(","8",",","8",")",")","plt.xlim","(","-2",",","2",")","plt.ylim","(","-2",",","2",")","plt.plot","(","samples","[",":",",","0","]",",","samples","[",":",",","1","]",",","'",".","'",")","plt.show","(",")","``","`","#","#","#","#","References","[","1","]",":","David","J.","Earl",",","Michael","W.","Deem","Parallel","Tempering",":","Theory",",","Applications",",","and","New","Perspectives","https",":","\/\/arxiv.org\/abs\/physics\/0508111","[","2","]",":","David","A.","Kofke","On","the","acceptance","probability","of","replica-exchange","Monte","Carlo","trials",".","J.","of","Chem",".","Phys",".","Vol",".","117","No",".","5","."]
code: def is_calibrated(self):
      return True
  
imports: import collections
  import warnings
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import broadcast_util
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.internal import unnest
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.mcmc.internal import util
label_desc: Use image node_id 12 for calling the ReplicaExchangeMC obj's underlying member method code with example usage: obj.is_calibrated() and returns: True