class_docstr: The experimental NeuTra kernel.
  
  Warning: This kernel is experimental. Default arguments and their
  interpreations will very likely change.
  
  This kernel does not work in TF1 graph mode.
  
  This implements a transition kernel that implements the NeuTra MCMC
  [(Hoffman et al., 2019)][1]. It operates by learning a neural pre-conditioner
  (bijector) inside the `bootstrap_results` method, and then using it inside
  `one_step` to accelerate HMC. The same bijector is used to initialize the
  state: i.e. the `current_state` as passed to `tfp.mcmc.sample_chain` is
  ignored except to extract the number of chains to sample in parallel.
  
  This kernel performs step-size adaptation and picks an automatic trajectory
  length for its internal HMC kernel.
  
  If your problem has constrained support, specify this support via the
  `unconstraining_bijector` argument. This argument is interpreted such that the
  image of the forward transformation of that bijector matches the support of
  your problem. E.g. if one of your random variables is positive, you could use
  the `tfb.Softplus` bijector.
  
  Since this is still experimental, we provide some facilities to debug the
  bijector training via the `train_debug_fn`, where you can monitor the training
  progress. In practice, it may be prudent to run the training process multiple
  times to verify that the variational approximation is stable. If it is not,
  you can attempt to increase the expressiveness of the bijector via the
  `trainable_bijector_fn` argument. Currently the bijector operates by
  flattening the entire problem state into one vector. This is seamless to the
  user, but is important to remember when designing a bijector. The flattening
  will likely be removed in the future.
  
  `_flattened_variational_distribution` returns the (flattened) distribution
  resulting from a standard normal being pushed through the bijector. This is
  useful for debugging, as this can be compared to the MCMC chain and your prior
  preconception of what the target distribution should look like.
  
  Additionally, you can examine the `inner_results.transformed_state` inside the
  kernel results. By construction, for a well-fitted bijector this should
  resemble a standard normal. This kernel assumes that you have a well-fitted
  bijector, so if samples in the transformed space do not look normal, then this
  kernel is not operating efficiently.
  
  ### Examples
  
  Sampling from a multivariate log-normal distribution.
  
  ```python
  
  target_dist = tfd.MultivariateNormalTriL(scale_tril=[[1., 0.], [2, 1.]])
  target_dist = tfb.Exp()(target_dist)
  
  num_chains = 64
  state_shape = 2
  num_burnin_steps = 1000
  
  kernel = NeuTra(
      target_log_prob_fn=target_dist.log_prob,
      state_shape=state_shape,
      num_step_size_adaptation_steps=int(0.8 * num_burnin_steps),
      unconstraining_bijector=tfb.Exp())
  
  chain = tfp.mcmc.sample_chain(
      num_results=1000,
      num_burnin_steps=1000,
      current_state=tf.ones([num_chains, 2]),
      kernel=kernel,
      trace_fn=None)
  ```
  
  #### References
  
  [1]: Hoffman, M., Sountsov, P., Dillon, J. V., Langmore, I., Tran, D., &
       Vasudevan, S. (2019). NeuTra-lizing Bad Geometry in Hamiltonian Monte
       Carlo Using Neural Transport. http://arxiv.org/abs/1903.03704
class_docstr_tok: ["The","experimental","NeuTra","kernel",".","Warning",":","This","kernel","is","experimental",".","Default","arguments","and","their","interpreations","will","very","likely","change",".","This","kernel","does","not","work","in","TF1","graph","mode",".","This","implements","a","transition","kernel","that","implements","the","NeuTra","MCMC","[","(","Hoffman","et","al.",",","2019",")","]","[","1","]",".","It","operates","by","learning","a","neural","pre-conditioner","(","bijector",")","inside","the","`","bootstrap_results","`","method",",","and","then","using","it","inside","`","one_step","`","to","accelerate","HMC",".","The","same","bijector","is","used","to","initialize","the","state",":","i.e",".","the","`","current_state","`","as","passed","to","`","tfp.mcmc.sample_chain","`","is","ignored","except","to","extract","the","number","of","chains","to","sample","in","parallel",".","This","kernel","performs","step-size","adaptation","and","picks","an","automatic","trajectory","length","for","its","internal","HMC","kernel",".","If","your","problem","has","constrained","support",",","specify","this","support","via","the","`","unconstraining_bijector","`","argument",".","This","argument","is","interpreted","such","that","the","image","of","the","forward","transformation","of","that","bijector","matches","the","support","of","your","problem",".","E.g",".","if","one","of","your","random","variables","is","positive",",","you","could","use","the","`","tfb.Softplus","`","bijector",".","Since","this","is","still","experimental",",","we","provide","some","facilities","to","debug","the","bijector","training","via","the","`","train_debug_fn","`",",","where","you","can","monitor","the","training","progress",".","In","practice",",","it","may","be","prudent","to","run","the","training","process","multiple","times","to","verify","that","the","variational","approximation","is","stable",".","If","it","is","not",",","you","can","attempt","to","increase","the","expressiveness","of","the","bijector","via","the","`","trainable_bijector_fn","`","argument",".","Currently","the","bijector","operates","by","flattening","the","entire","problem","state","into","one","vector",".","This","is","seamless","to","the","user",",","but","is","important","to","remember","when","designing","a","bijector",".","The","flattening","will","likely","be","removed","in","the","future",".","`","_flattened_variational_distribution","`","returns","the","(","flattened",")","distribution","resulting","from","a","standard","normal","being","pushed","through","the","bijector",".","This","is","useful","for","debugging",",","as","this","can","be","compared","to","the","MCMC","chain","and","your","prior","preconception","of","what","the","target","distribution","should","look","like",".","Additionally",",","you","can","examine","the","`","inner_results.transformed_state","`","inside","the","kernel","results",".","By","construction",",","for","a","well-fitted","bijector","this","should","resemble","a","standard","normal",".","This","kernel","assumes","that","you","have","a","well-fitted","bijector",",","so","if","samples","in","the","transformed","space","do","not","look","normal",",","then","this","kernel","is","not","operating","efficiently",".","#","#","#","Examples","Sampling","from","a","multivariate","log-normal","distribution",".","``","`","python","target_dist","=","tfd.MultivariateNormalTriL","(","scale_tril=","[","[","1.",",","0",".","]",",","[","2",",","1",".","]","]",")","target_dist","=","tfb.Exp","(",")","(","target_dist",")","num_chains","=","64","state_shape","=","2","num_burnin_steps","=","1000","kernel","=","NeuTra","(","target_log_prob_fn=target_dist.log_prob",",","state_shape=state_shape",",","num_step_size_adaptation_steps=int","(","0.8","*","num_burnin_steps",")",",","unconstraining_bijector=tfb.Exp","(",")",")","chain","=","tfp.mcmc.sample_chain","(","num_results=1000",",","num_burnin_steps=1000",",","current_state=tf.ones","(","[","num_chains",",","2","]",")",",","kernel=kernel",",","trace_fn=None",")","``","`","#","#","#","#","References","[","1","]",":","Hoffman",",","M.",",","Sountsov",",","P.",",","Dillon",",","J.","V.",",","Langmore",",","I.",",","Tran",",","D.",",","&","Vasudevan",",","S.","(","2019",")",".","NeuTra-lizing","Bad","Geometry","in","Hamiltonian","Monte","Carlo","Using","Neural","Transport",".","http",":","\/\/arxiv.org\/abs\/1903.03704"]
code: def _flattened_variational_distribution(self):
      base = tfd.MultivariateNormalDiag(
          loc=tf.zeros(self._total_event_size, dtype=self._dtype)
      )
      return self._bijector(base)
  
imports: import numpy
  import tensorflow.compat.v1
  import tensorflow_probability
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import unnest
label_desc: Use image node_id 15 for calling the NeuTra obj's underlying member method code with example usage: obj._flattened_variational_distribution() and returns: self