class_docstr: TransformedTransitionKernel applies a bijector to the MCMC's state space.
  
  The `TransformedTransitionKernel` `TransitionKernel` enables fitting
  a `tfp.bijectors.Bijector` which serves to decorrelate the Markov chain Monte
  Carlo (MCMC) event dimensions thus making the chain mix faster. This is
  particularly useful when the geometry of the target distribution is
  unfavorable. In such cases it may take many evaluations of the
  `target_log_prob_fn` for the chain to mix between faraway states.
  
  The idea of training an affine function to decorrelate chain event dims was
  presented in [Parno and Marzouk (2014)][1]. Used in conjunction with the
  `HamiltonianMonteCarlo` `TransitionKernel`, the [Parno and Marzouk (2014)][1]
  idea is an instance of Riemannian manifold HMC [(Girolami and Calderhead,
  2011)][2].
  
  The `TransformedTransitionKernel` enables arbitrary bijective transformations
  of arbitrary `TransitionKernel`s, e.g., one could use bijectors
  `tfp.bijectors.ScaleMatvecTriL`, `tfp.bijectors.RealNVP`, etc. with transition
  kernels `tfp.mcmc.HamiltonianMonteCarlo`, `tfp.mcmc.RandomWalkMetropolis`,
  etc.
  
  ### Transforming nested kernels
  
  `TransformedTransitionKernel` can operate on multiply nested kernels, as in
  the following example:
  
  ```python
  tfp.mcmc.TransformedTransitionKernel(
    inner_kernel=tfp.mcmc.SimpleStepSizeAdaptation(
      inner_kernel=tfp.mcmc.HamiltonianMonteCarlo(
        ... # doesn't matter
      ),
      num_adaptation_steps=9)
    bijector=tfb.Identity()))
  ```
  
  Upon construction, `TransformedTransitionKernel` searches the given
  `inner_kernel` and the "stack" of nested kernels in any `inner_kernel`
  fields thereof until it finds one with a field called `target_log_prob_fn`,
  and replaces this with the transformed function. If no
  `inner_kernel` has such a target log prob a `ValueError` is raised.
  
  #### Mathematical Details
  
  `TransformedTransitionKernel` enables Markov chains which operate in
  "unconstrained space." Since we interpret the bijector as mapping
  "unconstrained space" to "user space", this means that the MCMC transformed
  `target_log_prob` is:
  
  ```python
  target_log_prob(bij.forward(x)) + bij.forward_log_det_jacobian(x)
  ```
  
  Recall that `tfp.distributions.TransformedDistribution` uses the `inverse` to
  compute its `log_prob`. Despite this difference, the use of `forward` in
  `TransformedTransitionKernel` is perfectly consistent with
  `TransformedDistribution` following the TFP convention of "sampling" being
  what defines semantics. The apparent difference is because
  `TransformedDistribution.log_prob` is derived from a user provided
  distribution while in `TransformedTransitionKernel` samples are derived from
  `target_log_prob_fn`. That is, in `TransformedDistribution` we do:
  
  ```python
  x ~ NoiseDistribution()
  y = bij.forward(x)
  log_prob_y = NoiseDistribution().log_prob(bij.inverse(y))
               + bij.inverse_log_det_jacobian(y)
  ```
  
  yet in `TransformedTransitionKernel` we do:
  
  ```python
  x ~ MCMC()
  y = bij.forward(x)
  log_prob_y = log_prob(y) + bij.forward_log_det_jacobian(x)
  ```
  
  In other words (and in general), `tfp.mcmc` is derived from a `log_prob`
  which what induces a *seeming* direction convention change. Aside from TFP
  convention, that Bijectors should adhere to "sample first" semantics is
  important because it mitigates pervasive necessity of `tfp.bijectors.Invert`
  in user code.
  
  #### Examples
  
  ##### RealNVP + HamiltonianMonteCarlo
  
  Note: this example is only meant to illustrate how to wire up a
  `TransformedTransitionKernel`. As it is this won't work well because:
  * a 1-layer RealNVP is a pretty weak density model, since it can't change the
  density of the masked dimensions
  * we're not actually training the bijector to do anything useful.
  
  ```python
  import tensorflow as tf
  import tensorflow_probability as tfp
  tfd = tfp.distributions
  tfb = tfp.bijectors
  
  def make_likelihood(true_variances):
    return tfd.MultivariateNormalDiag(
        scale_diag=tf.sqrt(true_variances))
  
  dims = 10
  dtype = np.float32
  true_variances = tf.linspace(dtype(1), dtype(3), dims)
  likelihood = make_likelihood(true_variances)
  
  realnvp_hmc = tfp.mcmc.TransformedTransitionKernel(
      inner_kernel=tfp.mcmc.HamiltonianMonteCarlo(
        target_log_prob_fn=likelihood.log_prob,
        step_size=0.5,
        num_leapfrog_steps=2),
      bijector=tfb.RealNVP(
        num_masked=2,
        shift_and_log_scale_fn=tfb.real_nvp_default_template(
            hidden_layers=[512, 512])))
  
  states, kernel_results = tfp.mcmc.sample_chain(
      num_results=1000,
      current_state=tf.zeros(dims),
      kernel=realnvp_hmc,
      num_burnin_steps=500)
  
  # Compute sample stats.
  sample_mean = tf.reduce_mean(states, axis=0)
  sample_var = tf.reduce_mean(
      tf.squared_difference(states, sample_mean),
      axis=0)
  ```
  
  #### References
  
  [1]: Matthew Parno and Youssef Marzouk. Transport map accelerated Markov chain
       Monte Carlo. _arXiv preprint arXiv:1412.5492_, 2014.
       https://arxiv.org/abs/1412.5492
  
  [2]: Mark Girolami and Ben Calderhead. Riemann manifold langevin and
       hamiltonian monte carlo methods. In _Journal of the Royal Statistical
       Society_, 2011. https://doi.org/10.1111/j.1467-9868.2010.00765.x
class_docstr_tok: ["TransformedTransitionKernel","applies","a","bijector","to","the","MCMC","'s","state","space",".","The","`","TransformedTransitionKernel","`","`","TransitionKernel","`","enables","fitting","a","`","tfp.bijectors.Bijector","`","which","serves","to","decorrelate","the","Markov","chain","Monte","Carlo","(","MCMC",")","event","dimensions","thus","making","the","chain","mix","faster",".","This","is","particularly","useful","when","the","geometry","of","the","target","distribution","is","unfavorable",".","In","such","cases","it","may","take","many","evaluations","of","the","`","target_log_prob_fn","`","for","the","chain","to","mix","between","faraway","states",".","The","idea","of","training","an","affine","function","to","decorrelate","chain","event","dims","was","presented","in","[","Parno","and","Marzouk","(","2014",")","]","[","1","]",".","Used","in","conjunction","with","the","`","HamiltonianMonteCarlo","`","`","TransitionKernel","`",",","the","[","Parno","and","Marzouk","(","2014",")","]","[","1","]","idea","is","an","instance","of","Riemannian","manifold","HMC","[","(","Girolami","and","Calderhead",",","2011",")","]","[","2","]",".","The","`","TransformedTransitionKernel","`","enables","arbitrary","bijective","transformations","of","arbitrary","`","TransitionKernel","`","s",",","e.g.",",","one","could","use","bijectors","`","tfp.bijectors.ScaleMatvecTriL","`",",","`","tfp.bijectors.RealNVP","`",",","etc",".","with","transition","kernels","`","tfp.mcmc.HamiltonianMonteCarlo","`",",","`","tfp.mcmc.RandomWalkMetropolis","`",",","etc",".","#","#","#","Transforming","nested","kernels","`","TransformedTransitionKernel","`","can","operate","on","multiply","nested","kernels",",","as","in","the","following","example",":","``","`","python","tfp.mcmc.TransformedTransitionKernel","(","inner_kernel=tfp.mcmc.SimpleStepSizeAdaptation","(","inner_kernel=tfp.mcmc.HamiltonianMonteCarlo","(","...","#","does","n't","matter",")",",","num_adaptation_steps=9",")","bijector=tfb.Identity","(",")",")",")","``","`","Upon","construction",",","`","TransformedTransitionKernel","`","searches","the","given","`","inner_kernel","`","and","the","``","stack","''","of","nested","kernels","in","any","`","inner_kernel","`","fields","thereof","until","it","finds","one","with","a","field","called","`","target_log_prob_fn","`",",","and","replaces","this","with","the","transformed","function",".","If","no","`","inner_kernel","`","has","such","a","target","log","prob","a","`","ValueError","`","is","raised",".","#","#","#","#","Mathematical","Details","`","TransformedTransitionKernel","`","enables","Markov","chains","which","operate","in","''","unconstrained","space",".","''","Since","we","interpret","the","bijector","as","mapping","''","unconstrained","space","''","to","``","user","space","''",",","this","means","that","the","MCMC","transformed","`","target_log_prob","`","is",":","``","`","python","target_log_prob","(","bij.forward","(","x",")",")","+","bij.forward_log_det_jacobian","(","x",")","``","`","Recall","that","`","tfp.distributions.TransformedDistribution","`","uses","the","`","inverse","`","to","compute","its","`","log_prob","`",".","Despite","this","difference",",","the","use","of","`","forward","`","in","`","TransformedTransitionKernel","`","is","perfectly","consistent","with","`","TransformedDistribution","`","following","the","TFP","convention","of","``","sampling","''","being","what","defines","semantics",".","The","apparent","difference","is","because","`","TransformedDistribution.log_prob","`","is","derived","from","a","user","provided","distribution","while","in","`","TransformedTransitionKernel","`","samples","are","derived","from","`","target_log_prob_fn","`",".","That","is",",","in","`","TransformedDistribution","`","we","do",":","``","`","python","x","~","NoiseDistribution","(",")","y","=","bij.forward","(","x",")","log_prob_y","=","NoiseDistribution","(",")",".log_prob","(","bij.inverse","(","y",")",")","+","bij.inverse_log_det_jacobian","(","y",")","``","`","yet","in","`","TransformedTransitionKernel","`","we","do",":","``","`","python","x","~","MCMC","(",")","y","=","bij.forward","(","x",")","log_prob_y","=","log_prob","(","y",")","+","bij.forward_log_det_jacobian","(","x",")","``","`","In","other","words","(","and","in","general",")",",","`","tfp.mcmc","`","is","derived","from","a","`","log_prob","`","which","what","induces","a","*","seeming","*","direction","convention","change",".","Aside","from","TFP","convention",",","that","Bijectors","should","adhere","to","``","sample","first","''","semantics","is","important","because","it","mitigates","pervasive","necessity","of","`","tfp.bijectors.Invert","`","in","user","code",".","#","#","#","#","Examples","#","#","#","#","#","RealNVP","+","HamiltonianMonteCarlo","Note",":","this","example","is","only","meant","to","illustrate","how","to","wire","up","a","`","TransformedTransitionKernel","`",".","As","it","is","this","wo","n't","work","well","because",":","*","a","1-layer","RealNVP","is","a","pretty","weak","density","model",",","since","it","ca","n't","change","the","density","of","the","masked","dimensions","*","we","'re","not","actually","training","the","bijector","to","do","anything","useful",".","``","`","python","import","tensorflow","as","tf","import","tensorflow_probability","as","tfp","tfd","=","tfp.distributions","tfb","=","tfp.bijectors","def","make_likelihood","(","true_variances",")",":","return","tfd.MultivariateNormalDiag","(","scale_diag=tf.sqrt","(","true_variances",")",")","dims","=","10","dtype","=","np.float32","true_variances","=","tf.linspace","(","dtype","(","1",")",",","dtype","(","3",")",",","dims",")","likelihood","=","make_likelihood","(","true_variances",")","realnvp_hmc","=","tfp.mcmc.TransformedTransitionKernel","(","inner_kernel=tfp.mcmc.HamiltonianMonteCarlo","(","target_log_prob_fn=likelihood.log_prob",",","step_size=0.5",",","num_leapfrog_steps=2",")",",","bijector=tfb.RealNVP","(","num_masked=2",",","shift_and_log_scale_fn=tfb.real_nvp_default_template","(","hidden_layers=","[","512",",","512","]",")",")",")","states",",","kernel_results","=","tfp.mcmc.sample_chain","(","num_results=1000",",","current_state=tf.zeros","(","dims",")",",","kernel=realnvp_hmc",",","num_burnin_steps=500",")","#","Compute","sample","stats",".","sample_mean","=","tf.reduce_mean","(","states",",","axis=0",")","sample_var","=","tf.reduce_mean","(","tf.squared_difference","(","states",",","sample_mean",")",",","axis=0",")","``","`","#","#","#","#","References","[","1","]",":","Matthew","Parno","and","Youssef","Marzouk",".","Transport","map","accelerated","Markov","chain","Monte","Carlo",".","_arXiv","preprint","arXiv:1412.5492_",",","2014.","https",":","\/\/arxiv.org\/abs\/1412.5492","[","2","]",":","Mark","Girolami","and","Ben","Calderhead",".","Riemann","manifold","langevin","and","hamiltonian","monte","carlo","methods",".","In","_Journal","of","the","Royal","Statistical","Society_",",","2011.","https",":","\/\/doi.org\/10.1111\/j.1467-9868.2010.00765.x"]
code: def experimental_with_shard_axes(self, shard_axis_names):
      return self.copy(
          inner_kernel=self.inner_kernel.experimental_with_shard_axes(
              shard_axis_names
          )
      )
  
imports: import collections
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.mcmc.internal import util
  from tensorflow.python.util import deprecation
label_desc: Use image node_id 10 for calling the TransformedTransitionKernel obj's underlying member method code with example usage: obj.experimental_with_shard_axes(shard_axis_names) and returns: self