code: def sample_chain_with_burnin(
      num_results,
      current_state,
      previous_kernel_results=None,
      kernel=None,
      num_burnin_steps=0,
      num_steps_between_results=0,
      trace_fn=_trace_current_state,
      parallel_iterations=10,
      seed=None,
      name=None,
  ):
      """Implements Markov chain Monte Carlo via repeated `TransitionKernel` steps.
  
      This function samples from a Markov chain at `current_state` whose
      stationary distribution is governed by the supplied `TransitionKernel`
      instance (`kernel`).
  
      This function can sample from multiple chains, in parallel. (Whether or not
      there are multiple chains is dictated by the `kernel`.)
  
      The `current_state` can be represented as a single `Tensor` or a `list` of
      `Tensors` which collectively represent the current state.
  
      Since MCMC states are correlated, it is sometimes desirable to produce
      additional intermediate states, and then discard them, ending up with a set of
      states with decreased autocorrelation.  See [Owen (2017)][1]. Such 'thinning'
      is made possible by setting `num_steps_between_results > 0`. The chain then
      takes `num_steps_between_results` extra steps between the steps that make it
      into the results. The extra steps are never materialized, and thus do not
      increase memory requirements.
  
      In addition to returning the chain state, this function supports tracing of
      auxiliary variables used by the kernel. The traced values are selected by
      specifying `trace_fn`. By default, all chain states but no kernel results are
      traced.
  
      Args:
        num_results: Integer number of Markov chain draws.
        current_state: `Tensor` or Python `list` of `Tensor`s representing the
          current state(s) of the Markov chain(s).
        previous_kernel_results: A `Tensor` or a nested collection of `Tensor`s
          representing internal calculations made within the previous call to this
          function (or as returned by `bootstrap_results`).
        kernel: An instance of `tfp.mcmc.TransitionKernel` which implements one step
          of the Markov chain.
        num_burnin_steps: Integer number of chain steps to take before starting to
          collect results.
          Default value: 0 (i.e., no burn-in).
        num_steps_between_results: Integer number of chain steps between collecting
          a result. Only one out of every `num_steps_between_samples + 1` steps is
          included in the returned results.  The number of returned chain states is
          still equal to `num_results`.  Default value: 0 (i.e., no thinning).
        trace_fn: A callable that takes in the current chain state and the previous
          kernel results and return a `Tensor` or a nested collection of `Tensor`s
          that is then traced along with the chain state.
        parallel_iterations: The number of iterations allowed to run in parallel. It
          must be a positive integer. See `tf.while_loop` for more details.
        seed: PRNG seed; see `tfp.random.sanitize_seed` for details.
        name: Python `str` name prefixed to Ops created by this function.
          Default value: `None` (i.e.,
          'experimental_mcmc_sample_chain_with_burnin').
  
      Returns:
        result: A `RunKernelResults` instance containing information about the
          sampling run.  Main field is `trace`, the history of outputs of
          `trace_fn`.  See `RunKernelResults` for contents of other fields.
  
      #### References
  
      [1]: Art B. Owen. Statistically efficient thinning of a Markov chain sampler.
           _Technical Report_, 2017.
           http://statweb.stanford.edu/~owen/reports/bestthinning.pdf
      """
      with tf.name_scope(
          name or "experimental_mcmc_sample_chain_with_burnin"
      ):
          if not kernel.is_calibrated:
              warnings.warn(
                  "supplied `TransitionKernel` is not calibrated. Markov "
                  "chain may not converge to intended target distribution."
              )
  
          if trace_fn is None:
              trace_fn = lambda *args: ()
  
          burnin_seed, sampling_seed = samplers.split_seed(seed, n=2)
  
          # Burn-in run
          chain_state, kr = step.step_kernel(
              num_steps=num_burnin_steps,
              current_state=current_state,
              previous_kernel_results=previous_kernel_results,
              kernel=kernel,
              return_final_kernel_results=True,
              parallel_iterations=parallel_iterations,
              seed=burnin_seed,
              name="burnin",
          )
  
          thinning_k = thinning_kernel.ThinningKernel(
              kernel, num_steps_to_skip=num_steps_between_results
          )
  
          # ThinningKernel doesn't wrap the kernel_results structure, so we don't need
          # any of the usual munging.
          results = sample.sample_chain(
              num_results=num_results,
              current_state=chain_state,
              previous_kernel_results=kr,
              kernel=thinning_k,
              trace_fn=trace_fn,
              parallel_iterations=parallel_iterations,
              seed=sampling_seed,
              name="sampling",
          )
  
          del results.resume_kwargs["reducer"]
          del results.resume_kwargs["previous_reducer_state"]
          return results
  
file_path: probability/tensorflow_probability/python/experimental/mcmc/sample_fold.py
imports: import warnings
  import tensorflow.compat.v2
  from tensorflow_probability.python.experimental.mcmc import sample
  from tensorflow_probability.python.experimental.mcmc import sample_discarding_kernel
  from tensorflow_probability.python.experimental.mcmc import step
  from tensorflow_probability.python.experimental.mcmc import thinning_kernel
  from tensorflow_probability.python.experimental.mcmc import with_reductions
  from tensorflow_probability.python.internal import samplers
  from tensorflow.python.util import nest
label_desc: Use image node_id 3 for calling a global function with example usage: sample_chain_with_burnin(num_results, current_state, previous_kernel_results, kernel, num_burnin_steps, num_steps_between_results, trace_fn, parallel_iterations, seed, name) and returns: results