code: def window_tune_nuts_sampling(
      target_log_prob,
      prior_samples,
      constraining_bijectors=None,
      init_state=None,
      num_samples=500,
      nchains=4,
      init_nchains=1,
      target_accept_prob=0.8,
      max_tree_depth=9,
      use_scaled_init=True,
      tuning_window_schedule=(75, 25, 25, 25, 25, 25, 50),
      use_wide_window_expanding_mode=True,
      seed=None,
      parallel_iterations=10,
      jit_compile=True,
      use_input_signature=True,
      reduce_retracing=False,
  ):
      """Sample from a density with NUTS and an expanding window tuning scheme.
  
      This function implements a turnkey MCMC sampling routine using NUTS and an
      expanding window tuning strategy similar to Stan [1]. It learns a pre-
      conditioner that scales and rotates the target distribution using a series of
      expanding windows - either in number of samples (same as in Stan,
      use_wide_window_expanding_mode=False) or in number of batches/chains
      (use_wide_window_expanding_mode=True).
  
      Currently, the function uses `prior_samples` to initialize MCMC chains
      uniformly at random between -1 and 1 scaled by the prior standard deviation
      (i.e., [-prior_std, prior_std]). The scaling is ignored if `use_scaled_init`
      is set to False. Alternatively, user can input the `init_state` directly.
  
      Currently, the tuning and sampling routine is run in Python, with each block
      of the tuning epoch (window 1, 2, and 3 in Stan [1]) run with two @tf.function
      compiled functions. The user can control the compilation options using the
      kwargs `jit_compile`, `use_input_signature`, and
      `reduce_retracing`.  Setting all to True would compile to XLA and
      potentially avoid the small overhead of function recompilation (note that it
      is not yet the case in XLA right now). It is not yet clear whether doing it
      this way is better than just wrapping the full inference routine in
      tf.function with XLA.
  
      Internally, this calls `_sample_posterior`, which assumes a real-valued target
      density function and takes a Tensor with shape=(batch * dimension) as input.
      The tuning routine is a memory-less (i.e., no warm-up samples are saved) MCMC
      sampling with number of samples specified by a list-like
      `tuning_window_schedule`.
  
      Args:
        target_log_prob: 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.
        prior_samples: Nested structure of `Tensor`s, each of shape `[batches,
          latent_part_event_shape]` and should be sample from the prior. They are
          used to generate an initial chain position if `init_state` is not
          supplied.
        constraining_bijectors: `tfp.distributions.Bijector` or list of
          `tfp.distributions.Bijector`s. These bijectors use `forward` to map the
          state on the real space to the constrained state expected by
          `target_log_prob`.
        init_state: (Optional) `Tensor` or Python `list` of `Tensor`s representing
          the initial state(s) of the Markov chain(s).
        num_samples: Integer number of the Markov chain draws after tuning.
        nchains: Integer number of the Markov chains after tuning.
        init_nchains: Integer number of the Markov chains in the first phase of
          tuning.
        target_accept_prob: Floating point scalar `Tensor`. Target acceptance
          probability for step size adaptation.
        max_tree_depth: Maximum depth of the tree implicitly built by NUTS. See
          `tfp.mcmc.NoUTurnSampler` for more details
        use_scaled_init: Boolean. If `True`, generate initial state within [-1, 1]
          scaled by prior sample standard deviation in the unconstrained real space.
          This kwarg is ignored if `init_state` is not None
        tuning_window_schedule: List-like sequence of integers that specify the
          tuning schedule. Each integer number specifies the number of MCMC samples
          within a single warm-up window. The first and the last window tunes the
          step size (a scalar) only, while the intermediate windows tune both step
          size and the pre-conditioner. Moreover, the intermediate windows double
          the number of samples taken: for example, the default schedule (75,
            25, 25, 25, 25, 25, 50) actually means it will take (75, 25 * 2**0, 25 *
            2**1, 25 * 2**2, 25 * 2**3, 25 * 2**4, 50) samples.
        use_wide_window_expanding_mode: Boolean. Default to `True` that we double
          the number of chains from the previous stage for the intermediate windows.
          See `tuning_window_schedule` kwarg for more details.
        seed: Python integer to seed the random number generator.
        parallel_iterations: The number of iterations allowed to run in parallel.
          It must be a positive integer. See `tf.while_loop` for more details.
          Note that if you set the seed to have deterministic output you should
          also set `parallel_iterations` to 1.
        jit_compile: kwarg pass to tf.function decorator. If True, the
          function is always compiled by XLA.
        use_input_signature: If True, generate an input_signature kwarg to pass to
          tf.function decorator.
        reduce_retracing: kwarg pass to tf.function decorator. When True,
          tf.function may generate fewer, graphs that are less specialized on input
          shapes.
  
      Returns:
        posterior_samples: A `Tensor` or Python list of `Tensor`s representing the
          posterior MCMC samples after tuning. It has the same structure as
          `prior_samples` but with the leading shape being (num_samples * nchains)
        diagnostic: A list of `Tensor` representing the diagnostics from NUTS:
          `target_log_prob`, `leapfrogs_taken`, `has_divergence`, `energy`,
          `log_accept_ratio`, `reach_max_depth`, `step_size.
        conditioning_bijector: A tfp bijector that scales and rotates the target
          density function in latent unconstrained space as determined by
          adaptation.
  
      ### Examples
  
      Sampling from a multivariate Student-T distribution.
  
      ```python
      DTYPE = np.float32
  
      nd = 50
      concentration = 1.
  
      prior_dist = tfd.Sample(tfd.Normal(tf.constant(0., DTYPE), 100.), nd)
  
      mu = tf.cast(np.linspace(-100, 100, nd), dtype=DTYPE)
      sigma = tf.cast(np.exp(np.linspace(-1, 1.5, nd)), dtype=DTYPE)
      corr_tril = tfd.CholeskyLKJ(
          dimension=nd, concentration=concentration).sample()
      scale_tril = tf.linalg.matmul(tf.linalg.diag(sigma), corr_tril)
      target_dist = tfd.MultivariateStudentTLinearOperator(
          df=5., loc=mu, scale=tf.linalg.LinearOperatorLowerTriangular(scale_tril))
  
      target_log_prob = lambda *x: (
          prior_dist.log_prob(*x) + target_dist.log_prob(*x))
  
      (
          [mcmc_samples], diagnostic, conditioning_bijector
      ) = window_tune_nuts_sampling(target_log_prob, [prior_dist.sample(2000)])
  
      loc_conditioner, scale_conditioner = conditioning_bijector.trainable_variables
  
      _, ax = plt.subplots(1, 2, figsize=(10, 5))
      ax[0].plot(mu, loc_conditioner.numpy(), 'o', label='conditioner mean')
      ax[0].plot(mu, tf.reduce_mean(
          mcmc_samples, axis=[0, 1]), 'o', label='estimated mean')
      ax[0].legend()
  
      sigma_sim = target_dist._stddev()
      ax[1].plot(sigma_sim, scale_conditioner.numpy(), 'o', label='conditioner std')
      ax[1].plot(sigma_sim, tf.math.reduce_std(
          mcmc_samples, axis=[0, 1]), 'o', label='estimated std');
      ax[1].legend()
  
      ax[0].plot([min(mu), max(mu)], [min(mu), max(mu)])
      ax[1].plot([min(sigma_sim), max(sigma_sim)], [min(sigma_sim), max(sigma_sim)])
      ```
  
      #### References
  
      [1]: Stan Reference Manual.
      https://mc-stan.org/docs/2_23/reference-manual/hmc-algorithm-parameters.html#automatic-parameter-tuning
      """
  
      log_prob_val = target_log_prob(*prior_samples)
      log_prob_rank = ps.rank(log_prob_val)
      assert log_prob_rank == 1
  
      if constraining_bijectors is not None:
          target_log_prob_unconstrained = make_transformed_log_prob(
              target_log_prob,
              constraining_bijectors,
              direction="forward",
              enable_bijector_caching=False,
          )
          # constrain to unconstrain
          inverse_transform = make_transform_fn(
              constraining_bijectors, "inverse"
          )
          # unconstrain to constrain
          forward_transform = make_transform_fn(
              constraining_bijectors, "forward"
          )
      else:
          target_log_prob_unconstrained = target_log_prob
          inverse_transform = lambda x: x
          forward_transform = lambda y: y
  
      prior_samples_unconstrained = inverse_transform(prior_samples)
      init_state_unconstrained = None
  
      # If the input to target_log_prob_fn is a nested structure of Tensors, we
      # flatten and concatenate them into a 1D vector so that it is easier to work
      # with in mass matrix adaptation.
      if tf.nest.is_nested(prior_samples_unconstrained):
          free_rv_event_shape = [
              x.shape[log_prob_rank:] for x in prior_samples
          ]
          flat_event_splits = [
              s.num_elements() for s in free_rv_event_shape
          ]
  
          # TODO(b/158878248): replace the two function below with `tfb.Split`.
          def split_and_reshape(x):
              assertions = []
              message = "Input must have at least one dimension."
              if tensorshape_util.rank(x.shape) is not None:
                  if tensorshape_util.rank(x.shape) == 0:
                      raise ValueError(message)
              else:
                  assertions.append(
                      assert_util.assert_rank_at_least(
                          x, 1, message=message
                      )
                  )
              with tf.control_dependencies(assertions):
                  x = tf.nest.pack_sequence_as(
                      free_rv_event_shape,
                      tf.split(x, flat_event_splits, axis=-1),
                  )
  
                  def _reshape_map_part(part, event_shape):
                      static_rank = tf.get_static_value(
                          ps.rank_from_shape(event_shape)
                      )
                      if static_rank == 1:
                          return part
                      new_shape = ps.concat(
                          [ps.shape(part)[:-1], event_shape], axis=-1
                      )
                      return tf.reshape(
                          part, ps.cast(new_shape, tf.int32)
                      )
  
                  x = tf.nest.map_structure(
                      _reshape_map_part, x, free_rv_event_shape
                  )
              return x
  
          def concat_list_event(x):
              def handle_part(x, shape):
                  if (
                      len(shape) == 0
                  ):  # pylint: disable=g-explicit-length-test
                      return x[..., tf.newaxis]
                  return tf.reshape(
                      x, list(x.shape)[: -len(shape)] + [-1]
                  )
  
              flat_parts = [
                  handle_part(v, s)
                  for v, s in zip(x, free_rv_event_shape)
              ]
              return tf.concat(flat_parts, axis=-1)
  
          def target_log_prob_unconstrained_concated(x):
              x = split_and_reshape(x)
              return target_log_prob_unconstrained(*x)
  
          prior_samples_unconstrained_concated = concat_list_event(
              prior_samples_unconstrained
          )
          if init_state is not None:
              init_state_unconstrained = concat_list_event(
                  inverse_transform(init_state)
              )
      else:
          target_log_prob_unconstrained_concated = (
              target_log_prob_unconstrained
          )
          prior_samples_unconstrained_concated = (
              prior_samples_unconstrained
          )
          split_and_reshape = lambda x: x
          if init_state is not None:
              init_state_unconstrained = inverse_transform(init_state)
  
      (
          nuts_samples,
          diagnostic,
          conditioning_bijector,
      ) = _sample_posterior(
          target_log_prob_unconstrained_concated,
          prior_samples_unconstrained_concated,
          init_state=init_state_unconstrained,
          num_samples=num_samples,
          nchains=nchains,
          init_nchains=init_nchains,
          target_accept_prob=target_accept_prob,
          max_tree_depth=max_tree_depth,
          use_scaled_init=use_scaled_init,
          tuning_window_schedule=tuning_window_schedule,
          use_wide_window_expanding_mode=use_wide_window_expanding_mode,
          seed=seed,
          parallel_iterations=parallel_iterations,
          jit_compile=jit_compile,
          use_input_signature=use_input_signature,
          reduce_retracing=reduce_retracing,
      )
      return (
          forward_transform(split_and_reshape(nuts_samples)),
          diagnostic,
          conditioning_bijector,
      )
  
file_path: probability/discussion/turnkey_inference_candidate/window_tune_nuts_sampling.py
imports: from absl import logging
  import tensorflow.compat.v2
  import tensorflow_probability
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.mcmc.transformed_kernel import make_transform_fn
  from tensorflow_probability.python.mcmc.transformed_kernel import make_transformed_log_prob
label_desc: Use image node_id 2 for calling a global function with example usage: window_tune_nuts_sampling(target_log_prob, prior_samples, constraining_bijectors, init_state, num_samples, nchains, init_nchains, target_accept_prob, max_tree_depth, use_scaled_init, tuning_window_schedule, use_wide_window_expanding_mode, seed, parallel_iterations, jit_compile, use_input_signature, reduce_retracing) and returns: forward_transform,  diagnostic,  conditioning_bijector, x, tf, target_log_prob_unconstrained, tf, tf, x, part