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"]
docstr: Creates the kernel.
  
  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.
    state_shape: A Python list, or list of Python lists (or `TensorShape`s)
      that describes the shape of the problem state. Must be the same
      structure as `current_state`. All shapes must be fully defined.
    num_step_size_adaptation_steps: Number of steps to use for step size
      adaptation. See `tfp.mcmc.SimpleStepSizeAdaptation`.
    unconstraining_bijector: A bijector or list of bijectors that go from
      unconstrained space to the support of the corresponding random variable.
      Must be the same structure as `current_state`.
    trainable_bijector_fn: Creates a trainable, vector-valued event bijector.
      Must be a callable with signature: `(total_event_size, seed) ->
        bijector`, where `total_event_size` is the size of the event.
    learning_rate: Base learning rate to use for training the bijector.
      Internally, learning rate decay is used to stabilize learning.
    train_batch_size: Batch size to use for training the bijector.
    num_train_steps: Number of training steps to train the bijector.
    train_debug_fn: A callable with signature `(NeuTra, step, loss)` called
      for every training step. The first argument is this instance, and `step`
      is the current training step.
    seed: A seed for reproducibility.
docstr_tok: ["Creates","the","kernel",".","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",".","state_shape",":","A","Python","list",",","or","list","of","Python","lists","(","or","`","TensorShape","`","s",")","that","describes","the","shape","of","the","problem","state",".","Must","be","the","same","structure","as","`","current_state","`",".","All","shapes","must","be","fully","defined",".","num_step_size_adaptation_steps",":","Number","of","steps","to","use","for","step","size","adaptation",".","See","`","tfp.mcmc.SimpleStepSizeAdaptation","`",".","unconstraining_bijector",":","A","bijector","or","list","of","bijectors","that","go","from","unconstrained","space","to","the","support","of","the","corresponding","random","variable",".","Must","be","the","same","structure","as","`","current_state","`",".","trainable_bijector_fn",":","Creates","a","trainable",",","vector-valued","event","bijector",".","Must","be","a","callable","with","signature",":","`","(","total_event_size",",","seed",")","-",">","bijector","`",",","where","`","total_event_size","`","is","the","size","of","the","event",".","learning_rate",":","Base","learning","rate","to","use","for","training","the","bijector",".","Internally",",","learning","rate","decay","is","used","to","stabilize","learning",".","train_batch_size",":","Batch","size","to","use","for","training","the","bijector",".","num_train_steps",":","Number","of","training","steps","to","train","the","bijector",".","train_debug_fn",":","A","callable","with","signature","`","(","NeuTra",",","step",",","loss",")","`","called","for","every","training","step",".","The","first","argument","is","this","instance",",","and","`","step","`","is","the","current","training","step",".","seed",":","A","seed","for","reproducibility","."]
code: def __init__(
      self,
      target_log_prob_fn,
      state_shape,
      num_step_size_adaptation_steps,
      unconstraining_bijector=None,
      trainable_bijector_fn=make_iaf_stack,
      learning_rate=1e-2,
      train_batch_size=4096,
      num_train_steps=5000,
      train_debug_fn=None,
      seed=None,
  ):
      """Creates the kernel.
  
      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.
        state_shape: A Python list, or list of Python lists (or `TensorShape`s)
          that describes the shape of the problem state. Must be the same
          structure as `current_state`. All shapes must be fully defined.
        num_step_size_adaptation_steps: Number of steps to use for step size
          adaptation. See `tfp.mcmc.SimpleStepSizeAdaptation`.
        unconstraining_bijector: A bijector or list of bijectors that go from
          unconstrained space to the support of the corresponding random variable.
          Must be the same structure as `current_state`.
        trainable_bijector_fn: Creates a trainable, vector-valued event bijector.
          Must be a callable with signature: `(total_event_size, seed) ->
            bijector`, where `total_event_size` is the size of the event.
        learning_rate: Base learning rate to use for training the bijector.
          Internally, learning rate decay is used to stabilize learning.
        train_batch_size: Batch size to use for training the bijector.
        num_train_steps: Number of training steps to train the bijector.
        train_debug_fn: A callable with signature `(NeuTra, step, loss)` called
          for every training step. The first argument is this instance, and `step`
          is the current training step.
        seed: A seed for reproducibility.
      """
      self._parameters = dict(
          target_log_prob_fn=target_log_prob_fn,
          state_shape=state_shape,
          num_step_size_adaptation_steps=num_step_size_adaptation_steps,
          unconstraining_bijector=unconstraining_bijector,
          trainable_bijector_fn=trainable_bijector_fn,
          learning_rate=learning_rate,
          train_batch_size=train_batch_size,
          num_train_steps=num_train_steps,
          train_debug_fn=train_debug_fn,
          seed=seed,
      )
  
      self._state_shape = tf.nest.map_structure(
          tf.TensorShape, state_shape
      )
  
      if unconstraining_bijector is None:
          unconstraining_bijector = tf.nest.map_structure(
              lambda _: tfb.Identity(), self.state_shape
          )
      self._unconstraining_bijector = unconstraining_bijector
  
      def _make_reshaped_bijector(b, s):
          return tfb.Reshape(
              event_shape_in=s, event_shape_out=[s.num_elements()]
          )(b)(tfb.Reshape(event_shape_out=b.inverse_event_shape(s)))
  
      # This converts the `unconstraining_bijector` to work on flattened state
      # parts.
      reshaped_bijector = tf.nest.map_structure(
          _make_reshaped_bijector,
          unconstraining_bijector,
          self.state_shape,
      )
  
      blockwise_bijector = tfb.Blockwise(
          bijectors=tf.nest.flatten(reshaped_bijector),
          block_sizes=tf.nest.flatten(
              tf.nest.map_structure(
                  lambda s: s.num_elements(), self.state_shape
              )
          ),
      )
  
      dtype = dtype_util.common_dtype(
          [learning_rate], dtype_hint=tf.float32
      )
      self._dtype = dtype
  
      trainable_bijector = trainable_bijector_fn(
          self._total_event_size, seed=seed, dtype=dtype
      )
      self._trainable_bijector = trainable_bijector
  
      self._bijector = blockwise_bijector(trainable_bijector)
  
      def flattened_target_log_prob(flat_state):
          state = self._unflatten_state(flat_state)
          if isinstance(state, (list, tuple)):
              return target_log_prob_fn(*state)
          else:
              return target_log_prob_fn(state)
  
      self._flattened_target_log_prob_val = flattened_target_log_prob
  
      kernel = tfp.mcmc.HamiltonianMonteCarlo(
          target_log_prob_fn=flattened_target_log_prob,
          step_size=1.0,
          num_leapfrog_steps=self._num_leapfrog_steps(1.0),
      )
      kernel = tfp.mcmc.TransformedTransitionKernel(
          kernel, self._bijector
      )
      kernel = tfp.mcmc.SimpleStepSizeAdaptation(
          kernel,
          num_adaptation_steps=num_step_size_adaptation_steps,
          target_accept_prob=np.array(0.9, dtype.as_numpy_dtype),
      )
      self._kernel = kernel
  
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 1 to create a new NeuTra object from inherited base classes: tfp.mcmc with example: obj = NeuTra(target_log_prob_fn, state_shape, num_step_size_adaptation_steps, unconstraining_bijector, trainable_bijector_fn, learning_rate, train_batch_size, num_train_steps, train_debug_fn, seed)