Instruction:
Can you generate the documentation for the following Python 3 function
def _compute_log_acceptance_correction(current_momentums,
                                       proposed_momentums,
                                       independent_chain_ndims,
                                       name=None):
  """Helper to `kernel` which computes the log acceptance-correction.

  A sufficient but not necessary condition for the existence of a stationary
  distribution, `p(x)`, is "detailed balance", i.e.:

  ```none
  p(x'|x) p(x) = p(x|x') p(x')
  ```

  In the Metropolis-Hastings algorithm, a state is proposed according to
  `g(x'|x)` and accepted according to `a(x'|x)`, hence
  `p(x'|x) = g(x'|x) a(x'|x)`.

  Inserting this into the detailed balance equation implies:

  ```none
      g(x'|x) a(x'|x) p(x) = g(x|x') a(x|x') p(x')
  ==> a(x'|x) / a(x|x') = p(x') / p(x) [g(x|x') / g(x'|x)]    (*)
  ```

  One definition of `a(x'|x)` which satisfies (*) is:

  ```none
  a(x'|x) = min(1, p(x') / p(x) [g(x|x') / g(x'|x)])
  ```

  (To see that this satisfies (*), notice that under this definition only at
  most one `a(x'|x)` and `a(x|x') can be other than one.)

  We call the bracketed term the "acceptance correction".

  In the case of UncalibratedHMC, the log acceptance-correction is not the log
  proposal-ratio. UncalibratedHMC augments the state-space with momentum, z.
  Assuming a standard Gaussian distribution for momentums, the chain eventually
  converges to:

  ```none
  p([x, z]) propto= target_prob(x) exp(-0.5 z**2)
  ```

  Relating this back to Metropolis-Hastings parlance, for HMC we have:

  ```none
  p([x, z]) propto= target_prob(x) exp(-0.5 z**2)
  g([x, z] | [x', z']) = g([x', z'] | [x, z])
  ```

  In other words, the MH bracketed term is `1`. However, because we desire to
  use a general MH framework, we can place the momentum probability ratio inside
  the metropolis-correction factor thus getting an acceptance probability:

  ```none
                       target_prob(x')
  accept_prob(x'|x) = -----------------  [exp(-0.5 z**2) / exp(-0.5 z'**2)]
                       target_prob(x)
  ```

  (Note: we actually need to handle the kinetic energy change at each leapfrog
  step, but this is the idea.)

  Args:
    current_momentums: `Tensor` representing the value(s) of the current
      momentum(s) of the state (parts).
    proposed_momentums: `Tensor` representing the value(s) of the proposed
      momentum(s) of the state (parts).
    independent_chain_ndims: Scalar `int` `Tensor` representing the number of
      leftmost `Tensor` dimensions which index independent chains.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., 'compute_log_acceptance_correction').

  Returns:
    log_acceptance_correction: `Tensor` representing the `log`
      acceptance-correction.  (See docstring for mathematical definition.)
  """
  with tf.compat.v1.name_scope(
      name, 'compute_log_acceptance_correction',
      [independent_chain_ndims, current_momentums, proposed_momentums]):
    log_current_kinetic, log_proposed_kinetic = [], []
    for current_momentum, proposed_momentum in zip(
        current_momentums, proposed_momentums):
      axis = tf.range(independent_chain_ndims, tf.rank(current_momentum))
      log_current_kinetic.append(_log_sum_sq(current_momentum, axis))
      log_proposed_kinetic.append(_log_sum_sq(proposed_momentum, axis))
    current_kinetic = 0.5 * tf.exp(
        tf.reduce_logsumexp(
            input_tensor=tf.stack(log_current_kinetic, axis=-1), axis=-1))
    proposed_kinetic = 0.5 * tf.exp(
        tf.reduce_logsumexp(
            input_tensor=tf.stack(log_proposed_kinetic, axis=-1), axis=-1))
    return mcmc_util.safe_sum([current_kinetic, -proposed_kinetic])

Response: Helper function to compute the log acceptance correction.


Instruction:
How would you explain what the following Python 3 function does
def _log_sum_sq(x, axis=None):
  return tf.reduce_logsumexp(
      input_tensor=2. * tf.math.log(tf.abs(x)), axis=axis)

Response: Computes log sum of x.


Instruction:
How would you explain what the following Python 3 function does
def one_step(self, current_state, previous_kernel_results):
    previous_step_size_assign = (
        [] if self.step_size_update_fn is None
        else (previous_kernel_results.extra.step_size_assign
              if mcmc_util.is_list_like(
                  previous_kernel_results.extra.step_size_assign)
              else [previous_kernel_results.extra.step_size_assign]))

    with tf.control_dependencies(previous_step_size_assign):
      next_state, kernel_results = self._impl.one_step(
          current_state, previous_kernel_results)
      if self.step_size_update_fn is not None:
        step_size_assign = self.step_size_update_fn(  # pylint: disable=not-callable
            self.step_size, kernel_results)
        kernel_results = kernel_results._replace(
            extra=HamiltonianMonteCarloExtraKernelResults(
                step_size_assign=step_size_assign))
      return next_state, kernel_results

Response: Runs one iteration of Hamiltonian Monte Carlo.


Instruction:
Make a summary of the following Python 3 code
def bootstrap_results(self, init_state):
    kernel_results = self._impl.bootstrap_results(init_state)
    if self.step_size_update_fn is not None:
      step_size_assign = self.step_size_update_fn(self.step_size, None)  # pylint: disable=not-callable
      kernel_results = kernel_results._replace(
          extra=HamiltonianMonteCarloExtraKernelResults(
              step_size_assign=step_size_assign))
    return kernel_results

Response: Creates initial kernel_results using a supplied state.


Instruction:
Explain what the following Python 3 code does
def bayesian_resnet(input_shape,
                    num_classes=10,
                    kernel_posterior_scale_mean=-9.0,
                    kernel_posterior_scale_stddev=0.1,
                    kernel_posterior_scale_constraint=0.2):
  """Constructs a ResNet18 model.

  Args:
    input_shape: A `tuple` indicating the Tensor shape.
    num_classes: `int` representing the number of class labels.
    kernel_posterior_scale_mean: Python `int` number for the kernel
      posterior's scale (log variance) mean. The smaller the mean the closer
      is the initialization to a deterministic network.
    kernel_posterior_scale_stddev: Python `float` number for the initial kernel
      posterior's scale stddev.
      ```
      q(W|x) ~ N(mu, var),
      log_var ~ N(kernel_posterior_scale_mean, kernel_posterior_scale_stddev)
      ````
    kernel_posterior_scale_constraint: Python `float` number for the log value
      to constrain the log variance throughout training.
      i.e. log_var <= log(kernel_posterior_scale_constraint).

  Returns:
    tf.keras.Model.
  """

  filters = [64, 128, 256, 512]
  kernels = [3, 3, 3, 3]
  strides = [1, 2, 2, 2]

  def _untransformed_scale_constraint(t):
    return tf.clip_by_value(t, -1000,
                            tf.math.log(kernel_posterior_scale_constraint))

  kernel_posterior_fn = tfp.layers.default_mean_field_normal_fn(
      untransformed_scale_initializer=tf.compat.v1.initializers.random_normal(
          mean=kernel_posterior_scale_mean,
          stddev=kernel_posterior_scale_stddev),
      untransformed_scale_constraint=_untransformed_scale_constraint)

  image = tf.keras.layers.Input(shape=input_shape, dtype='float32')
  x = tfp.layers.Convolution2DFlipout(
      64,
      3,
      strides=1,
      padding='same',
      kernel_posterior_fn=kernel_posterior_fn)(image)

  for i in range(len(kernels)):
    x = _resnet_block(
        x,
        filters[i],
        kernels[i],
        strides[i],
        kernel_posterior_fn)

  x = tf.keras.layers.BatchNormalization()(x)
  x = tf.keras.layers.Activation('relu')(x)
  x = tf.keras.layers.AveragePooling2D(4, 1)(x)
  x = tf.keras.layers.Flatten()(x)

  x = tfp.layers.DenseFlipout(
      num_classes,
      kernel_posterior_fn=kernel_posterior_fn)(x)

  model = tf.keras.Model(inputs=image, outputs=x, name='resnet18')
  return model

Response: Constructs a ResNet18 model.


Instruction:
Given the following Python 3 function, write the documentation
def _resnet_block(x, filters, kernel, stride, kernel_posterior_fn):
  x = tf.keras.layers.BatchNormalization()(x)
  x = tf.keras.layers.Activation('relu')(x)

  if stride != 1 or filters != x.shape[1]:
    shortcut = _projection_shortcut(x, filters, stride, kernel_posterior_fn)
  else:
    shortcut = x

  x = tfp.layers.Convolution2DFlipout(
      filters,
      kernel,
      strides=stride,
      padding='same',
      kernel_posterior_fn=kernel_posterior_fn)(x)
  x = tf.keras.layers.BatchNormalization()(x)
  x = tf.keras.layers.Activation('relu')(x)

  x = tfp.layers.Convolution2DFlipout(
      filters,
      kernel,
      strides=1,
      padding='same',
      kernel_posterior_fn=kernel_posterior_fn)(x)
  x = tf.keras.layers.add([x, shortcut])
  return x

Response: Network block for ResNet.


Instruction:
Can you generate the documentation for the following Python 3 function
def make_encoder(activation, num_topics, layer_sizes):
  encoder_net = tf.keras.Sequential()
  for num_hidden_units in layer_sizes:
    encoder_net.add(
        tf.keras.layers.Dense(
            num_hidden_units,
            activation=activation,
            kernel_initializer=tf.compat.v1.glorot_normal_initializer()))
  encoder_net.add(
      tf.keras.layers.Dense(
          num_topics,
          activation=tf.nn.softplus,
          kernel_initializer=tf.compat.v1.glorot_normal_initializer()))

  def encoder(bag_of_words):
    net = _clip_dirichlet_parameters(encoder_net(bag_of_words))
    return tfd.Dirichlet(concentration=net,
                         name="topics_posterior")

  return encoder

Response: Create the encoder function.


Instruction:
Can you create a Python 3 function that
creates the decoder function.

Response:
def make_decoder(num_topics, num_words):
  """Create the decoder function.

  Args:
    num_topics: The number of topics.
    num_words: The number of words.

  Returns:
    decoder: A `callable` mapping a `Tensor` of encodings to a
      `tfd.Distribution` instance over words.
  """
  topics_words_logits = tf.compat.v1.get_variable(
      "topics_words_logits",
      shape=[num_topics, num_words],
      initializer=tf.compat.v1.glorot_normal_initializer())
  topics_words = tf.nn.softmax(topics_words_logits, axis=-1)

  def decoder(topics):
    word_probs = tf.matmul(topics, topics_words)
    # The observations are bag of words and therefore not one-hot. However,
    # log_prob of OneHotCategorical computes the probability correctly in
    # this case.
    return tfd.OneHotCategorical(probs=word_probs,
                                 name="bag_of_words")

  return decoder, topics_words


Instruction:
Can you implement a function in Python 3 that
creates the prior distribution.

Response:
def make_prior(num_topics, initial_value):
  """Create the prior distribution.

  Args:
    num_topics: Number of topics.
    initial_value: The starting value for the prior parameters.

  Returns:
    prior: A `callable` that returns a `tf.distribution.Distribution`
        instance, the prior distribution.
    prior_variables: A `list` of `Variable` objects, the trainable parameters
        of the prior.
  """
  def _softplus_inverse(x):
    return np.log(np.expm1(x))

  logit_concentration = tf.compat.v1.get_variable(
      "logit_concentration",
      shape=[1, num_topics],
      initializer=tf.compat.v1.initializers.constant(
          _softplus_inverse(initial_value)))
  concentration = _clip_dirichlet_parameters(
      tf.nn.softplus(logit_concentration))

  def prior():
    return tfd.Dirichlet(concentration=concentration,
                         name="topics_prior")

  prior_variables = [logit_concentration]

  return prior, prior_variables


Instruction:
Can you generate the documentation for the following Python 3 function
def model_fn(features, labels, mode, params, config):
  del labels, config

  encoder = make_encoder(params["activation"],
                         params["num_topics"],
                         params["layer_sizes"])
  decoder, topics_words = make_decoder(params["num_topics"],
                                       features.shape[1])
  prior, prior_variables = make_prior(params["num_topics"],
                                      params["prior_initial_value"])

  topics_prior = prior()
  alpha = topics_prior.concentration

  topics_posterior = encoder(features)
  topics = topics_posterior.sample()
  random_reconstruction = decoder(topics)

  reconstruction = random_reconstruction.log_prob(features)
  tf.compat.v1.summary.scalar("reconstruction",
                              tf.reduce_mean(input_tensor=reconstruction))

  # Compute the KL-divergence between two Dirichlets analytically.
  # The sampled KL does not work well for "sparse" distributions
  # (see Appendix D of [2]).
  kl = tfd.kl_divergence(topics_posterior, topics_prior)
  tf.compat.v1.summary.scalar("kl", tf.reduce_mean(input_tensor=kl))

  # Ensure that the KL is non-negative (up to a very small slack).
  # Negative KL can happen due to numerical instability.
  with tf.control_dependencies(
      [tf.compat.v1.assert_greater(kl, -1e-3, message="kl")]):
    kl = tf.identity(kl)

  elbo = reconstruction - kl
  avg_elbo = tf.reduce_mean(input_tensor=elbo)
  tf.compat.v1.summary.scalar("elbo", avg_elbo)
  loss = -avg_elbo

  # Perform variational inference by minimizing the -ELBO.
  global_step = tf.compat.v1.train.get_or_create_global_step()
  optimizer = tf.compat.v1.train.AdamOptimizer(params["learning_rate"])

  # This implements the "burn-in" for prior parameters (see Appendix D of [2]).
  # For the first prior_burn_in_steps steps they are fixed, and then trained
  # jointly with the other parameters.
  grads_and_vars = optimizer.compute_gradients(loss)
  grads_and_vars_except_prior = [
      x for x in grads_and_vars if x[1] not in prior_variables]

  def train_op_except_prior():
    return optimizer.apply_gradients(
        grads_and_vars_except_prior,
        global_step=global_step)

  def train_op_all():
    return optimizer.apply_gradients(
        grads_and_vars,
        global_step=global_step)

  train_op = tf.cond(
      pred=global_step < params["prior_burn_in_steps"],
      true_fn=train_op_except_prior,
      false_fn=train_op_all)

  # The perplexity is an exponent of the average negative ELBO per word.
  words_per_document = tf.reduce_sum(input_tensor=features, axis=1)
  log_perplexity = -elbo / words_per_document
  tf.compat.v1.summary.scalar(
      "perplexity", tf.exp(tf.reduce_mean(input_tensor=log_perplexity)))
  (log_perplexity_tensor,
   log_perplexity_update) = tf.compat.v1.metrics.mean(log_perplexity)
  perplexity_tensor = tf.exp(log_perplexity_tensor)

  # Obtain the topics summary. Implemented as a py_func for simplicity.
  topics = tf.compat.v1.py_func(
      functools.partial(get_topics_strings, vocabulary=params["vocabulary"]),
      [topics_words, alpha],
      tf.string,
      stateful=False)
  tf.compat.v1.summary.text("topics", topics)

  return tf.estimator.EstimatorSpec(
      mode=mode,
      loss=loss,
      train_op=train_op,
      eval_metric_ops={
          "elbo": tf.compat.v1.metrics.mean(elbo),
          "reconstruction": tf.compat.v1.metrics.mean(reconstruction),
          "kl": tf.compat.v1.metrics.mean(kl),
          "perplexity": (perplexity_tensor, log_perplexity_update),
          "topics": (topics, tf.no_op()),
      },
  )

Response: Builds the model function for use in an estimator.


Instruction:
Explain what the following Python 3 code does
def sample_chain(
    num_results,
    current_state,
    previous_kernel_results=None,
    kernel=None,
    num_burnin_steps=0,
    num_steps_between_results=0,
    trace_fn=lambda current_state, kernel_results: kernel_results,
    return_final_kernel_results=False,
    parallel_iterations=10,
    name=None,
):
  """Implements Markov chain Monte Carlo via repeated `TransitionKernel` steps.

  This function samples from an Markov chain at `current_state` and 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 (in calls to
  `sess.run`), and thus do not increase memory requirements.

  Warning: when setting a `seed` in the `kernel`, ensure that `sample_chain`'s
  `parallel_iterations=1`, otherwise results will not be reproducible.

  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 kernel results are traced but in the
  future the default will be changed to no results being traced, so plan
  accordingly. See below for some examples of this feature.

  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.
    return_final_kernel_results: If `True`, then the final kernel results are
      returned alongside the chain state and the trace specified by the
      `trace_fn`.
    parallel_iterations: The number of iterations allowed to run in parallel. It
      must be a positive integer. See `tf.while_loop` for more details.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., "mcmc_sample_chain").

  Returns:
    checkpointable_states_and_trace: if `return_final_kernel_results` is
      `True`. The return value is an instance of
      `CheckpointableStatesAndTrace`.
    all_states: if `return_final_kernel_results` is `False` and `trace_fn` is
      `None`. The return value is a `Tensor` or Python list of `Tensor`s
      representing the state(s) of the Markov chain(s) at each result step. Has
      same shape as input `current_state` but with a prepended
      `num_results`-size dimension.
    states_and_trace: if `return_final_kernel_results` is `False` and
      `trace_fn` is not `None`. The return value is an instance of
      `StatesAndTrace`.

  #### Examples

  ##### Sample from a diagonal-variance Gaussian.

  I.e.,

  ```none
  for i=1..n:
    x[i] ~ MultivariateNormal(loc=0, scale=diag(true_stddev))  # likelihood
  ```

  ```python
  import tensorflow as tf
  import tensorflow_probability as tfp
  tfd = tfp.distributions

  dims = 10
  true_stddev = np.sqrt(np.linspace(1., 3., dims))
  likelihood = tfd.MultivariateNormalDiag(loc=0., scale_diag=true_stddev)

  states = tfp.mcmc.sample_chain(
      num_results=1000,
      num_burnin_steps=500,
      current_state=tf.zeros(dims),
      kernel=tfp.mcmc.HamiltonianMonteCarlo(
        target_log_prob_fn=likelihood.log_prob,
        step_size=0.5,
        num_leapfrog_steps=2),
      trace_fn=None)

  sample_mean = tf.reduce_mean(states, axis=0)
  # ==> approx all zeros

  sample_stddev = tf.sqrt(tf.reduce_mean(
      tf.squared_difference(states, sample_mean),
      axis=0))
  # ==> approx equal true_stddev
  ```

  ##### Sampling from factor-analysis posteriors with known factors.

  I.e.,

  ```none
  # prior
  w ~ MultivariateNormal(loc=0, scale=eye(d))
  for i=1..n:
    # likelihood
    x[i] ~ Normal(loc=w^T F[i], scale=1)
  ```

  where `F` denotes factors.

  ```python
  import tensorflow as tf
  import tensorflow_probability as tfp
  tfd = tfp.distributions

  # Specify model.
  def make_prior(dims):
    return tfd.MultivariateNormalDiag(
        loc=tf.zeros(dims))

  def make_likelihood(weights, factors):
    return tfd.MultivariateNormalDiag(
        loc=tf.matmul(weights, factors, adjoint_b=True))

  def joint_log_prob(num_weights, factors, x, w):
    return (make_prior(num_weights).log_prob(w) +
            make_likelihood(w, factors).log_prob(x))

  def unnormalized_log_posterior(w):
    # Posterior is proportional to: `p(W, X=x | factors)`.
    return joint_log_prob(num_weights, factors, x, w)

  # Setup data.
  num_weights = 10 # == d
  num_factors = 40 # == n
  num_chains = 100

  weights = make_prior(num_weights).sample(1)
  factors = tf.random_normal([num_factors, num_weights])
  x = make_likelihood(weights, factors).sample()

  # Sample from Hamiltonian Monte Carlo Markov Chain.

  # Get `num_results` samples from `num_chains` independent chains.
  chains_states, kernels_results = tfp.mcmc.sample_chain(
      num_results=1000,
      num_burnin_steps=500,
      current_state=tf.zeros([num_chains, num_weights], name='init_weights'),
      kernel=tfp.mcmc.HamiltonianMonteCarlo(
        target_log_prob_fn=unnormalized_log_posterior,
        step_size=0.1,
        num_leapfrog_steps=2))

  # Compute sample stats.
  sample_mean = tf.reduce_mean(chains_states, axis=[0, 1])
  # ==> approx equal to weights

  sample_var = tf.reduce_mean(
      tf.squared_difference(chains_states, sample_mean),
      axis=[0, 1])
  # ==> less than 1
  ```

  ##### Custom tracing functions.

  ```python
  import tensorflow as tf
  import tensorflow_probability as tfp
  tfd = tfp.distributions

  likelihood = tfd.Normal(loc=0., scale=1.)

  def sample_chain(trace_fn):
    return tfp.mcmc.sample_chain(
      num_results=1000,
      num_burnin_steps=500,
      current_state=0.,
      kernel=tfp.mcmc.HamiltonianMonteCarlo(
        target_log_prob_fn=likelihood.log_prob,
        step_size=0.5,
        num_leapfrog_steps=2),
      trace_fn=trace_fn)

  def trace_log_accept_ratio(states, previous_kernel_results):
    return previous_kernel_results.log_accept_ratio

  def trace_everything(states, previous_kernel_results):
    return previous_kernel_results

  _, log_accept_ratio = sample_chain(trace_fn=trace_log_accept_ratio)
  _, kernel_results = sample_chain(trace_fn=trace_everything)

  acceptance_prob = tf.exp(tf.minimum(log_accept_ratio_, 0.))
  # Equivalent to, but more efficient than:
  acceptance_prob = tf.exp(tf.minimum(kernel_results.log_accept_ratio_, 0.))
  ```

  #### References

  [1]: Art B. Owen. Statistically efficient thinning of a Markov chain sampler.
       _Technical Report_, 2017.
       http://statweb.stanford.edu/~owen/reports/bestthinning.pdf
  """
  if not kernel.is_calibrated:
    warnings.warn("supplied `TransitionKernel` is not calibrated. Markov "
                  "chain may not converge to intended target distribution.")
  with tf.compat.v1.name_scope(
      name, "mcmc_sample_chain",
      [num_results, num_burnin_steps, num_steps_between_results]):
    num_results = tf.convert_to_tensor(
        value=num_results, dtype=tf.int32, name="num_results")
    num_burnin_steps = tf.convert_to_tensor(
        value=num_burnin_steps, dtype=tf.int32, name="num_burnin_steps")
    num_steps_between_results = tf.convert_to_tensor(
        value=num_steps_between_results,
        dtype=tf.int32,
        name="num_steps_between_results")
    current_state = tf.nest.map_structure(
        lambda x: tf.convert_to_tensor(value=x, name="current_state"),
        current_state)
    if previous_kernel_results is None:
      previous_kernel_results = kernel.bootstrap_results(current_state)

    if trace_fn is None:
      # It simplifies the logic to use a dummy function here.
      trace_fn = lambda *args: ()
      no_trace = True
    else:
      no_trace = False
    if trace_fn is sample_chain.__defaults__[4]:
      warnings.warn("Tracing all kernel results by default is deprecated. Set "
                    "the `trace_fn` argument to None (the future default "
                    "value) or an explicit callback that traces the values "
                    "you are interested in.")

    def _trace_scan_fn(state_and_results, num_steps):
      next_state, current_kernel_results = mcmc_util.smart_for_loop(
          loop_num_iter=num_steps,
          body_fn=kernel.one_step,
          initial_loop_vars=list(state_and_results),
          parallel_iterations=parallel_iterations)
      return next_state, current_kernel_results

    (_, final_kernel_results), (all_states, trace) = mcmc_util.trace_scan(
        loop_fn=_trace_scan_fn,
        initial_state=(current_state, previous_kernel_results),
        elems=tf.one_hot(
            indices=0,
            depth=num_results,
            on_value=1 + num_burnin_steps,
            off_value=1 + num_steps_between_results,
            dtype=tf.int32),
        # pylint: disable=g-long-lambda
        trace_fn=lambda state_and_results: (state_and_results[0],
                                            trace_fn(*state_and_results)),
        # pylint: enable=g-long-lambda
        parallel_iterations=parallel_iterations)

    if return_final_kernel_results:
      return CheckpointableStatesAndTrace(
          all_states=all_states,
          trace=trace,
          final_kernel_results=final_kernel_results)
    else:
      if no_trace:
        return all_states
      else:
        return StatesAndTrace(all_states=all_states, trace=trace)

Response: Implements Markov chain Monte Carlo via repeated TransitionKernel steps.


Instruction:
Explain what the following Python 3 code does
def trainable_gamma(shape, min_concentration=1e-3, min_scale=1e-5, name=None):
  with tf.compat.v1.variable_scope(None, default_name="trainable_gamma"):
    unconstrained_concentration = tf.compat.v1.get_variable(
        "unconstrained_concentration",
        shape,
        initializer=tf.compat.v1.initializers.random_normal(
            mean=0.5, stddev=0.1))
    unconstrained_scale = tf.compat.v1.get_variable(
        "unconstrained_scale",
        shape,
        initializer=tf.compat.v1.initializers.random_normal(stddev=0.1))
    concentration = tf.maximum(tf.nn.softplus(unconstrained_concentration),
                               min_concentration)
    rate = tf.maximum(1. / tf.nn.softplus(unconstrained_scale), 1. / min_scale)
    rv = ed.Gamma(concentration=concentration, rate=rate, name=name)
    return rv

Response: Learnable Gamma via concentration and scale parameterization.


Instruction:
Create a Python 3 function for
loading NIPS 2011 conference papers from CSV file.

Response:
def load_nips2011_papers(path):
  """Loads NIPS 2011 conference papers.

  The NIPS 1987-2015 data set is in the form of a 11,463 x 5,812 matrix of
  per-paper word counts, containing 11,463 words and 5,811 NIPS conference
  papers (Perrone et al., 2016). We subset to papers in 2011 and words appearing
  in at least two documents and having a total word count of at least 10.

  Built from the Observations Python package.

  Args:
    path: str.
      Path to directory which either stores file or otherwise file will
      be downloaded and extracted there. Filename is `NIPS_1987-2015.csv`.

  Returns:
    bag_of_words: np.ndarray of shape [num_documents, num_words]. Each element
      denotes the number of occurrences of a specific word in a specific
      document.
    words: List of strings, denoting the words for `bag_of_words`'s columns.
  """
  path = os.path.expanduser(path)
  filename = "NIPS_1987-2015.csv"
  filepath = os.path.join(path, filename)
  if not os.path.exists(filepath):
    url = ("https://archive.ics.uci.edu/ml/machine-learning-databases/"
           "00371/NIPS_1987-2015.csv")
    if not tf.io.gfile.exists(path):
      tf.io.gfile.makedirs(path)
    print("Downloading %s to %s" % (url, filepath))
    urllib.request.urlretrieve(url, filepath)

  with open(filepath) as f:
    iterator = csv.reader(f)
    documents = next(iterator)[1:]
    words = []
    x_train = []
    for row in iterator:
      words.append(row[0])
      x_train.append(row[1:])

  x_train = np.array(x_train, dtype=np.int)

  # Subset to documents in 2011 and words appearing in at least two documents
  # and have a total word count of at least 10.
  doc_idx = [i for i, document in enumerate(documents)
             if document.startswith("2011")]
  documents = [documents[doc] for doc in doc_idx]
  x_train = x_train[:, doc_idx]
  word_idx = np.logical_and(np.sum(x_train != 0, 1) >= 2,
                            np.sum(x_train, 1) >= 10)
  words = [word for word, idx in zip(words, word_idx) if idx]
  bag_of_words = x_train[word_idx, :].T
  return bag_of_words, words


Instruction:
Write a Python 3 function for
sharing init logic for amplitude and length_scale params.

Response:
def _init_params(self, amplitude, length_scale, validate_args):
    """Shared init logic for `amplitude` and `length_scale` params.

    Args:
      amplitude: `Tensor` (or convertible) or `None` to convert, validate.
      length_scale: `Tensor` (or convertible) or `None` to convert, validate.
      validate_args: If `True`, parameters are checked for validity despite
        possibly degrading runtime performance

    Returns:
      dtype: The common `DType` of the parameters.
    """
    dtype = util.maybe_get_common_dtype(
        [amplitude, length_scale])
    if amplitude is not None:
      amplitude = tf.convert_to_tensor(
          value=amplitude, name='amplitude', dtype=dtype)
    self._amplitude = _validate_arg_if_not_none(
        amplitude, tf.compat.v1.assert_positive, validate_args)
    if length_scale is not None:
      length_scale = tf.convert_to_tensor(
          value=length_scale, name='length_scale', dtype=dtype)
    self._length_scale = _validate_arg_if_not_none(
        length_scale, tf.compat.v1.assert_positive, validate_args)
    return dtype


Instruction:
Can you create a Python 3 function that
gets the KL function registered for classes a and b.

Response:
def _registered_kl(type_a, type_b):
  """Get the KL function registered for classes a and b."""
  hierarchy_a = tf_inspect.getmro(type_a)
  hierarchy_b = tf_inspect.getmro(type_b)
  dist_to_children = None
  kl_fn = None
  for mro_to_a, parent_a in enumerate(hierarchy_a):
    for mro_to_b, parent_b in enumerate(hierarchy_b):
      candidate_dist = mro_to_a + mro_to_b
      candidate_kl_fn = _DIVERGENCES.get((parent_a, parent_b), None)
      if not kl_fn or (candidate_kl_fn and candidate_dist < dist_to_children):
        dist_to_children = candidate_dist
        kl_fn = candidate_kl_fn
  return kl_fn


Instruction:
Create a Python 3 function for
returning the batchwise KL - divergence between two distributions.

Response:
def kl_divergence(distribution_a, distribution_b,
                  allow_nan_stats=True, name=None):
  """Get the KL-divergence KL(distribution_a || distribution_b).

  If there is no KL method registered specifically for `type(distribution_a)`
  and `type(distribution_b)`, then the class hierarchies of these types are
  searched.

  If one KL method is registered between any pairs of classes in these two
  parent hierarchies, it is used.

  If more than one such registered method exists, the method whose registered
  classes have the shortest sum MRO paths to the input types is used.

  If more than one such shortest path exists, the first method
  identified in the search is used (favoring a shorter MRO distance to
  `type(distribution_a)`).

  Args:
    distribution_a: The first distribution.
    distribution_b: The second distribution.
    allow_nan_stats: Python `bool`, default `True`. When `True`,
      statistics (e.g., mean, mode, variance) use the value "`NaN`" to
      indicate the result is undefined. When `False`, an exception is raised
      if one or more of the statistic's batch members are undefined.
    name: Python `str` name prefixed to Ops created by this class.

  Returns:
    A Tensor with the batchwise KL-divergence between `distribution_a`
    and `distribution_b`.

  Raises:
    NotImplementedError: If no KL method is defined for distribution types
      of `distribution_a` and `distribution_b`.
  """
  kl_fn = _registered_kl(type(distribution_a), type(distribution_b))
  if kl_fn is None:
    raise NotImplementedError(
        "No KL(distribution_a || distribution_b) registered for distribution_a "
        "type {} and distribution_b type {}".format(
            type(distribution_a).__name__, type(distribution_b).__name__))

  with tf.name_scope("KullbackLeibler"):
    kl_t = kl_fn(distribution_a, distribution_b, name=name)
    if allow_nan_stats:
      return kl_t

    # Check KL for NaNs
    kl_t = tf.identity(kl_t, name="kl")

    with tf.control_dependencies([
        tf.Assert(
            tf.logical_not(tf.reduce_any(input_tensor=tf.math.is_nan(kl_t))),
            [
                ("KL calculation between {} and {} returned NaN values "
                 "(and was called with allow_nan_stats=False). Values:".format(
                     distribution_a.name, distribution_b.name)),
                kl_t
            ])
    ]):
      return tf.identity(kl_t, name="checked_kl")


Instruction:
Can you write a function in Python 3 where it
computes the cross entropy of two distributions by P and Q.

Response:
def cross_entropy(ref, other,
                  allow_nan_stats=True, name=None):
  """Computes the (Shannon) cross entropy.

  Denote two distributions by `P` (`ref`) and `Q` (`other`). Assuming `P, Q`
  are absolutely continuous with respect to one another and permit densities
  `p(x) dr(x)` and `q(x) dr(x)`, (Shanon) cross entropy is defined as:

  ```none
  H[P, Q] = E_p[-log q(X)] = -int_F p(x) log q(x) dr(x)
  ```

  where `F` denotes the support of the random variable `X ~ P`.

  Args:
    ref: `tfd.Distribution` instance.
    other: `tfd.Distribution` instance.
    allow_nan_stats: Python `bool`, default `True`. When `True`,
      statistics (e.g., mean, mode, variance) use the value "`NaN`" to
      indicate the result is undefined. When `False`, an exception is raised
      if one or more of the statistic's batch members are undefined.
    name: Python `str` prepended to names of ops created by this function.

  Returns:
    cross_entropy: `ref.dtype` `Tensor` with shape `[B1, ..., Bn]`
      representing `n` different calculations of (Shanon) cross entropy.
  """
  with tf.name_scope(name or "cross_entropy"):
    return ref.entropy() + kl_divergence(
        ref, other, allow_nan_stats=allow_nan_stats)


Instruction:
Can you generate a brief explanation for the following Python 3 code
def read_image(filepath):
  im_bytes = tf.io.read_file(filepath)
  im = tf.image.decode_image(im_bytes, channels=CHANNELS)
  im = tf.image.convert_image_dtype(im, tf.float32)
  return im

Response: Reads an image tensor from a file.


Instruction:
Explain what the following Python 3 code does
def download_sprites():
  filepath = os.path.join(FLAGS.data_dir, DATA_SPRITES_DIR)
  if not tf.io.gfile.exists(filepath):
    if not tf.io.gfile.exists(FLAGS.data_dir):
      tf.io.gfile.makedirs(FLAGS.data_dir)
    zip_name = "{}.zip".format(filepath)
    urllib.request.urlretrieve(DATA_SPRITES_URL, zip_name)
    with zipfile.ZipFile(zip_name, "r") as zip_file:
      zip_file.extractall(FLAGS.data_dir)
    tf.io.gfile.remove(zip_name)
  return filepath

Response: Downloads the sprites data and returns the saved filepath.


Instruction:
Can you write a function in Python 3 where it
creates a character sprite from a set of attribute sprites.

Response:
def create_character(skin, hair, top, pants):
  """Creates a character sprite from a set of attribute sprites."""
  dtype = skin.dtype
  hair_mask = tf.cast(hair[..., -1:] <= 0, dtype)
  top_mask = tf.cast(top[..., -1:] <= 0, dtype)
  pants_mask = tf.cast(pants[..., -1:] <= 0, dtype)
  char = (skin * hair_mask) + hair
  char = (char * top_mask) + top
  char = (char * pants_mask) + pants
  return char


Instruction:
How would you code a function in Python 3 to
create a sequence tensor for the given character sprite.

Response:
def create_seq(character, action_metadata, direction, length=8, start=0):
  """Creates a sequence.

  Args:
    character: A character sprite tensor.
    action_metadata: An action metadata tuple.
    direction: An integer representing the direction, i.e., the row
      offset within each action group corresponding to a particular
      direction.
    length: Desired length of the sequence. If this is longer than
      the number of available frames, it will roll over to the
      beginning.
    start: Index of possible frames at which to start the sequence.

  Returns:
    A sequence tensor.
  """
  sprite_start = (action_metadata[0]+direction) * FRAME_SIZE
  sprite_end = (action_metadata[0]+direction+1) * FRAME_SIZE
  sprite_line = character[sprite_start:sprite_end, ...]

  # Extract 64x64 patches that are side-by-side in the sprite, and limit
  # to the actual number of frames for the given action.
  frames = tf.stack(tf.split(sprite_line, 13, axis=1))  # 13 is a hack
  frames = frames[0:action_metadata[1]]

  # Extract a slice of the desired length.
  # NOTE: Length could be longer than the number of frames, so tile as needed.
  frames = tf.roll(frames, shift=-start, axis=0)
  frames = tf.tile(frames, [2, 1, 1, 1])  # 2 is a hack
  frames = frames[:length]
  frames = tf.cast(frames, dtype=tf.float32)
  frames.set_shape([length, FRAME_SIZE, FRAME_SIZE, CHANNELS])
  return frames


Instruction:
Create a Python 3 function to
create a random sequence.

Response:
def create_random_seq(character, action_metadata, direction, length=8):
  """Creates a random sequence."""
  start = tf.random.uniform([], maxval=action_metadata[1], dtype=tf.int32)
  return create_seq(character, action_metadata, direction, length, start)


Instruction:
Implement a Python 3 function for
creating a tf. data. Dataset for the sprites dataset.

Response:
def create_sprites_dataset(characters, actions, directions, channels=3,
                           length=8, shuffle=False, fake_data=False):
  """Creates a tf.data pipeline for the sprites dataset.

  Args:
    characters: A list of (skin, hair, top, pants) tuples containing
      relative paths to the sprite png image for each attribute.
    actions: A list of Actions.
    directions: A list of Directions.
    channels: Number of image channels to yield.
    length: Desired length of the sequences.
    shuffle: Whether or not to shuffle the characters and sequences
      start frame.
    fake_data: Boolean for whether or not to yield synthetic data.

  Returns:
    A tf.data.Dataset yielding (seq, skin label index, hair label index,
    top label index, pants label index, action label index, skin label
    name, hair label_name, top label name, pants label name, action
    label name) tuples.
  """
  if fake_data:
    dummy_image = tf.random.normal([HEIGHT, WIDTH, CHANNELS])
  else:
    basedir = download_sprites()

  action_names = [action.name for action in actions]
  action_metadata = [(action.start_row, action.frames) for action in actions]

  direction_rows = [direction.row_offset for direction in directions]

  chars = tf.data.Dataset.from_tensor_slices(characters)
  act_names = tf.data.Dataset.from_tensor_slices(action_names).repeat()
  acts_metadata = tf.data.Dataset.from_tensor_slices(action_metadata).repeat()
  dir_rows = tf.data.Dataset.from_tensor_slices(direction_rows).repeat()

  if shuffle:
    chars = chars.shuffle(len(characters))

  dataset = tf.data.Dataset.zip((chars, act_names, acts_metadata, dir_rows))

  skin_table = tf.contrib.lookup.index_table_from_tensor(sorted(SKIN_COLORS))
  hair_table = tf.contrib.lookup.index_table_from_tensor(sorted(HAIRSTYLES))
  top_table = tf.contrib.lookup.index_table_from_tensor(sorted(TOPS))
  pants_table = tf.contrib.lookup.index_table_from_tensor(sorted(PANTS))
  action_table = tf.contrib.lookup.index_table_from_tensor(sorted(action_names))

  def process_example(attrs, act_name, act_metadata, dir_row_offset):
    """Processes a dataset row."""
    skin_name = attrs[0]
    hair_name = attrs[1]
    top_name = attrs[2]
    pants_name = attrs[3]

    if fake_data:
      char = dummy_image
    else:
      skin = read_image(basedir + os.sep + skin_name)
      hair = read_image(basedir + os.sep + hair_name)
      top = read_image(basedir + os.sep + top_name)
      pants = read_image(basedir + os.sep + pants_name)
      char = create_character(skin, hair, top, pants)

    if shuffle:
      seq = create_random_seq(char, act_metadata, dir_row_offset, length)
    else:
      seq = create_seq(char, act_metadata, dir_row_offset, length)
    seq = seq[..., :channels]  # limit output channels

    skin_idx = skin_table.lookup(skin_name)
    hair_idx = hair_table.lookup(hair_name)
    top_idx = top_table.lookup(top_name)
    pants_idx = pants_table.lookup(pants_name)
    act_idx = action_table.lookup(act_name)

    return (seq, skin_idx, hair_idx, top_idx, pants_idx, act_idx,
            skin_name, hair_name, top_name, pants_name, act_name)

  dataset = dataset.map(process_example)
  return dataset


Instruction:
Here you have a function in Python 3, explain what it does
def _maybe_validate_distributions(distributions, dtype_override, validate_args):
  assertions = []

  if not _is_iterable(distributions) or not distributions:
    raise ValueError('`distributions` must be a list of one or more '
                     'distributions.')

  if dtype_override is None:
    dts = [
        dtype_util.base_dtype(d.dtype)
        for d in distributions
        if d.dtype is not None
    ]
    if dts[1:] != dts[:-1]:
      raise TypeError('Distributions must have same dtype; found: {}.'.format(
          set(dtype_util.name(dt) for dt in dts)))

  # Validate event_ndims.
  for d in distributions:
    if tensorshape_util.rank(d.event_shape) is not None:
      if tensorshape_util.rank(d.event_shape) != 1:
        raise ValueError('`Distribution` must be vector variate, '
                         'found event nimds: {}.'.format(
                             tensorshape_util.rank(d.event_shape)))
    elif validate_args:
      assertions.append(
          assert_util.assert_equal(
              1, tf.size(input=d.event_shape_tensor()),
              message='`Distribution` must be vector variate.'))

  batch_shapes = [d.batch_shape for d in distributions]
  if all(tensorshape_util.is_fully_defined(b) for b in batch_shapes):
    if batch_shapes[1:] != batch_shapes[:-1]:
      raise ValueError('Distributions must have the same `batch_shape`; '
                       'found: {}.'.format(batch_shapes))
  elif validate_args:
    batch_shapes = [
        tensorshape_util.as_list(d.batch_shape)  # pylint: disable=g-complex-comprehension
        if tensorshape_util.is_fully_defined(d.batch_shape) else
        d.batch_shape_tensor() for d in distributions
    ]
    assertions.extend(
        assert_util.assert_equal(  # pylint: disable=g-complex-comprehension
            b1, b2,
            message='Distribution `batch_shape`s must be identical.')
        for b1, b2 in zip(batch_shapes[1:], batch_shapes[:-1]))

  return assertions

Response: Checks that the given distributions satisfy all assumptions.


Instruction:
Can you generate the documentation for the following Python 3 function
def _kl_blockwise_blockwise(b0, b1, name=None):
  if len(b0.distributions) != len(b1.distributions):
    raise ValueError(
        'Can only compute KL divergence between Blockwise distributions with '
        'the same number of component distributions.')

  # We also need to check that the event shapes match for each one.
  b0_event_sizes = [_event_size(d) for d in b0.distributions]
  b1_event_sizes = [_event_size(d) for d in b1.distributions]

  assertions = []
  message = ('Can only compute KL divergence between Blockwise distributions '
             'with the same pairwise event shapes.')

  if (all(isinstance(event_size, int) for event_size in b0_event_sizes) and
      all(isinstance(event_size, int) for event_size in b1_event_sizes)):
    if b0_event_sizes != b1_event_sizes:
      raise ValueError(message)
  else:
    if b0.validate_args or b1.validate_args:
      assertions.extend(
          assert_util.assert_equal(  # pylint: disable=g-complex-comprehension
              e1, e2, message=message)
          for e1, e2 in zip(b0_event_sizes, b1_event_sizes))

  with tf.name_scope(name or 'kl_blockwise_blockwise'):
    with tf.control_dependencies(assertions):
      return sum([
          kullback_leibler.kl_divergence(d1, d2) for d1, d2 in zip(
              b0.distributions, b1.distributions)])

Response: Calculate the batched KL divergence between two Blockwise distributions.


Instruction:
Here you have a function in Python 3, explain what it does
def _kl_half_normal_half_normal(a, b, name=None):
  with tf.name_scope(name or "kl_half_normal_half_normal"):
    # Consistent with
    # http://www.mast.queensu.ca/~communications/Papers/gil-msc11.pdf, page 119
    return (tf.math.log(b.scale) - tf.math.log(a.scale) +
            (a.scale**2 - b.scale**2) / (2 * b.scale**2))

Response: Calculate the batched KL divergence KL with a and b HalfNormal distribution.


Instruction:
Write a Python 3 script to
flatten a list of kernels which may contain _SumKernel instances.

Response:
def _flatten_summand_list(kernels):
  """Flatten a list of kernels which may contain _SumKernel instances.

  Args:
    kernels: Python list of `PositiveSemidefiniteKernel` instances

  Returns:
    Python list containing the elements of kernels, with any _SumKernel
    instances replaced by their `kernels` property contents.
  """
  flattened = []
  for k in kernels:
    if isinstance(k, _SumKernel):
      flattened += k.kernels
    else:
      flattened.append(k)
  return flattened


Instruction:
Here you have a function in Python 3, explain what it does
def _flatten_multiplicand_list(kernels):
  flattened = []
  for k in kernels:
    if isinstance(k, _ProductKernel):
      flattened += k.kernels
    else:
      flattened.append(k)
  return flattened

Response: Flatten a list of kernels which may contain _ProductKernel instances.


Instruction:
How would you explain what the following Python 3 function does
def build_input_pipeline(x_train, x_test, y_train, y_test,
                         batch_size, valid_size):
  """Build an Iterator switching between train and heldout data."""

  x_train = x_train.astype("float32")
  x_test = x_test.astype("float32")

  x_train /= 255
  x_test /= 255

  y_train = y_train.flatten()
  y_test = y_test.flatten()

  if FLAGS.subtract_pixel_mean:
    x_train_mean = np.mean(x_train, axis=0)
    x_train -= x_train_mean
    x_test -= x_train_mean

  print("x_train shape:" + str(x_train.shape))
  print(str(x_train.shape[0]) + " train samples")
  print(str(x_test.shape[0]) + " test samples")

  # Build an iterator over training batches.
  training_dataset = tf.data.Dataset.from_tensor_slices(
      (x_train, np.int32(y_train)))
  training_batches = training_dataset.shuffle(
      50000, reshuffle_each_iteration=True).repeat().batch(batch_size)
  training_iterator = tf.compat.v1.data.make_one_shot_iterator(training_batches)

  # Build a iterator over the heldout set with batch_size=heldout_size,
  # i.e., return the entire heldout set as a constant.
  heldout_dataset = tf.data.Dataset.from_tensor_slices(
      (x_test, np.int32(y_test)))
  heldout_batches = heldout_dataset.repeat().batch(valid_size)
  heldout_iterator = tf.compat.v1.data.make_one_shot_iterator(heldout_batches)

  # Combine these into a feedable iterator that can switch between training
  # and validation inputs.
  handle = tf.compat.v1.placeholder(tf.string, shape=[])
  feedable_iterator = tf.compat.v1.data.Iterator.from_string_handle(
      handle, training_batches.output_types, training_batches.output_shapes)
  images, labels = feedable_iterator.get_next()

  return images, labels, handle, training_iterator, heldout_iterator

Response: Builds a pipeline that can be used to switch between train and heldout data.


Instruction:
Make a summary of the following Python 3 code
def build_fake_data():
  num_examples = 10
  x_train = np.random.rand(num_examples, *IMAGE_SHAPE).astype(np.float32)
  y_train = np.random.permutation(np.arange(num_examples)).astype(np.int32)
  x_test = np.random.rand(num_examples, *IMAGE_SHAPE).astype(np.float32)
  y_test = np.random.permutation(np.arange(num_examples)).astype(np.int32)
  return (x_train, y_train), (x_test, y_test)

Response: Build fake CIFAR10 - style data for unit testing.


Instruction:
How would you code a function in Python 3 to
count the number of occurrences of each value in an integer array.

Response:
def count_integers(arr,
                   weights=None,
                   minlength=None,
                   maxlength=None,
                   axis=None,
                   dtype=tf.int32,
                   name=None):
  """Counts the number of occurrences of each value in an integer array `arr`.

  Works like `tf.math.bincount`, but provides an `axis` kwarg that specifies
  dimensions to reduce over.  With
    `~axis = [i for i in range(arr.ndim) if i not in axis]`,
  this function returns a `Tensor` of shape `[K] + arr.shape[~axis]`.

  If `minlength` and `maxlength` are not given, `K = tf.reduce_max(arr) + 1`
  if `arr` is non-empty, and 0 otherwise.
  If `weights` are non-None, then index `i` of the output stores the sum of the
  value in `weights` at each index where the corresponding value in `arr` is
  `i`.

  Args:
    arr: An `int32` `Tensor` of non-negative values.
    weights: If non-None, must be the same shape as arr. For each value in
      `arr`, the bin will be incremented by the corresponding weight instead of
      1.
    minlength: If given, ensures the output has length at least `minlength`,
      padding with zeros at the end if necessary.
    maxlength: If given, skips values in `arr` that are equal or greater than
      `maxlength`, ensuring that the output has length at most `maxlength`.
    axis: A `0-D` or `1-D` `int32` `Tensor` (with static values) designating
      dimensions in `arr` to reduce over.
      `Default value:` `None`, meaning reduce over all dimensions.
    dtype: If `weights` is None, determines the type of the output bins.
    name: A name scope for the associated operations (optional).

  Returns:
    A vector with the same dtype as `weights` or the given `dtype`. The bin
    values.
  """
  with tf.compat.v1.name_scope(
      name, 'count_integers', values=[arr, weights, minlength, maxlength,
                                      axis]):
    if axis is None:
      return tf.math.bincount(
          arr,
          weights=weights,
          minlength=minlength,
          maxlength=maxlength,
          dtype=dtype)

    arr = tf.convert_to_tensor(value=arr, dtype=tf.int32, name='arr')
    arr_ndims = _get_static_ndims(arr, expect_static=True)

    axis = _make_static_axis_non_negative_list(axis, arr_ndims)

    # ~axis from docstring.  Dims in arr that are not in axis.
    not_axis = sorted(set(range(arr_ndims)).difference(axis))

    # If we're reducing over everything, just use standard bincount.
    if not not_axis:
      return tf.math.bincount(
          arr,
          weights=weights,
          minlength=minlength,
          maxlength=maxlength,
          dtype=dtype)

    # Move dims in ~axis to the left, so we can tf.map_fn bincount over them,
    # Producing counts for every index I in ~axis.
    # Thus, flat_arr is not totally flat, it just has the dims in ~axis
    # flattened.
    flat_arr = _move_dims_to_flat_end(arr, not_axis, arr_ndims, right_end=False)

    # tf.map_fn over dim 0.
    if weights is None:

      def one_bincount(arr_slice):
        return tf.math.bincount(
            arr_slice,
            weights=None,
            minlength=minlength,
            maxlength=maxlength,
            dtype=dtype)

      flat_counts = tf.map_fn(one_bincount, elems=flat_arr, dtype=dtype)
    else:
      weights = tf.convert_to_tensor(value=weights, name='weights')
      _get_static_ndims(weights, expect_static=True, expect_ndims=arr_ndims)
      flat_weights = _move_dims_to_flat_end(
          weights, not_axis, arr_ndims, right_end=False)

      def one_bincount(arr_and_weights_slices):
        arr_slice, weights_slice = arr_and_weights_slices
        return tf.math.bincount(
            arr_slice,
            weights=weights_slice,
            minlength=minlength,
            maxlength=maxlength,
            dtype=dtype)

      flat_counts = tf.map_fn(
          one_bincount, elems=[flat_arr, flat_weights], dtype=weights.dtype)

    # flat_counts.shape = [prod(~axis), K], because map_fn stacked on axis 0.
    # bincount needs to have the K bins in axis 0, so transpose...
    flat_counts_t = tf.transpose(a=flat_counts, perm=[1, 0])

    # Throw in this assert, to ensure shape assumptions are correct.
    _get_static_ndims(flat_counts_t, expect_ndims=2, expect_static=True)

    # not_axis_shape = arr.shape[~axis]
    not_axis_shape = tf.gather(tf.shape(input=arr), indices=not_axis)

    # The first index of flat_counts_t indexes bins 0,..,K-1, the rest are ~axis
    out_shape = tf.concat([[-1], not_axis_shape], axis=0)

    return tf.reshape(flat_counts_t, out_shape)


Instruction:
Write a Python 3 function that can
return a list of bin values into discrete intervals.

Response:
def find_bins(x,
              edges,
              extend_lower_interval=False,
              extend_upper_interval=False,
              dtype=None,
              name=None):
  """Bin values into discrete intervals.

  Given `edges = [c0, ..., cK]`, defining intervals
  `I0 = [c0, c1)`, `I1 = [c1, c2)`, ..., `I_{K-1} = [c_{K-1}, cK]`,
  This function returns `bins`, such that:
  `edges[bins[i]] <= x[i] < edges[bins[i] + 1]`.

  Args:
    x:  Numeric `N-D` `Tensor` with `N > 0`.
    edges:  `Tensor` of same `dtype` as `x`.  The first dimension indexes edges
      of intervals.  Must either be `1-D` or have
      `x.shape[1:] == edges.shape[1:]`.  If `rank(edges) > 1`, `edges[k]`
      designates a shape `edges.shape[1:]` `Tensor` of bin edges for the
      corresponding dimensions of `x`.
    extend_lower_interval:  Python `bool`.  If `True`, extend the lowest
      interval `I0` to `(-inf, c1]`.
    extend_upper_interval:  Python `bool`.  If `True`, extend the upper
      interval `I_{K-1}` to `[c_{K-1}, +inf)`.
    dtype: The output type (`int32` or `int64`). `Default value:` `x.dtype`.
      This effects the output values when `x` is below/above the intervals,
      which will be `-1/K+1` for `int` types and `NaN` for `float`s.
      At indices where `x` is `NaN`, the output values will be `0` for `int`
      types and `NaN` for floats.
    name:  A Python string name to prepend to created ops. Default: 'find_bins'

  Returns:
    bins: `Tensor` with same `shape` as `x` and `dtype`.
      Has whole number values.  `bins[i] = k` means the `x[i]` falls into the
      `kth` bin, ie, `edges[bins[i]] <= x[i] < edges[bins[i] + 1]`.

  Raises:
    ValueError:  If `edges.shape[0]` is determined to be less than 2.

  #### Examples

  Cut a `1-D` array

  ```python
  x = [0., 5., 6., 10., 20.]
  edges = [0., 5., 10.]
  tfp.stats.find_bins(x, edges)
  ==> [0., 0., 1., 1., np.nan]
  ```

  Cut `x` into its deciles

  ```python
  x = tf.random_uniform(shape=(100, 200))
  decile_edges = tfp.stats.quantiles(x, num_quantiles=10)
  bins = tfp.stats.find_bins(x, edges=decile_edges)
  bins.shape
  ==> (100, 200)
  tf.reduce_mean(bins == 0.)
  ==> approximately 0.1
  tf.reduce_mean(bins == 1.)
  ==> approximately 0.1
  ```

  """
  # TFP users may be surprised to see the "action" in the leftmost dim of
  # edges, rather than the rightmost (event) dim.  Why?
  # 1. Most likely you created edges by getting quantiles over samples, and
  #    quantile/percentile return these edges in the leftmost (sample) dim.
  # 2. Say you have event_shape = [5], then we expect the bin will be different
  #    for all 5 events, so the index of the bin should not be in the event dim.
  with tf.compat.v1.name_scope(
      name, default_name='find_bins', values=[x, edges]):
    in_type = dtype_util.common_dtype([x, edges],
                                      preferred_dtype=tf.float32)
    edges = tf.convert_to_tensor(value=edges, name='edges', dtype=in_type)
    x = tf.convert_to_tensor(value=x, name='x', dtype=in_type)

    if (tf.compat.dimension_value(edges.shape[0]) is not None and
        tf.compat.dimension_value(edges.shape[0]) < 2):
      raise ValueError(
          'First dimension of `edges` must have length > 1 to index 1 or '
          'more bin. Found: {}'.format(edges.shape))

    flattening_x = edges.shape.ndims == 1 and x.shape.ndims > 1

    if flattening_x:
      x_orig_shape = tf.shape(input=x)
      x = tf.reshape(x, [-1])

    if dtype is None:
      dtype = in_type
    dtype = tf.as_dtype(dtype)

    # Move first dims into the rightmost.
    x_permed = distribution_util.rotate_transpose(x, shift=-1)
    edges_permed = distribution_util.rotate_transpose(edges, shift=-1)

    # If...
    #   x_permed = [0, 1, 6., 10]
    #   edges = [0, 5, 10.]
    #   ==> almost_output = [0, 1, 2, 2]
    searchsorted_type = dtype if dtype in [tf.int32, tf.int64] else None
    almost_output_permed = tf.searchsorted(
        sorted_sequence=edges_permed,
        values=x_permed,
        side='right',
        out_type=searchsorted_type)
    # Move the rightmost dims back to the leftmost.
    almost_output = tf.cast(
        distribution_util.rotate_transpose(almost_output_permed, shift=1),
        dtype)

    # In above example, we want [0, 0, 1, 1], so correct this here.
    bins = tf.clip_by_value(almost_output - 1, tf.cast(0, dtype),
                            tf.cast(tf.shape(input=edges)[0] - 2, dtype))

    if not extend_lower_interval:
      low_fill = np.nan if dtype.is_floating else -1
      bins = tf.where(x < tf.expand_dims(edges[0], 0),
                      tf.fill(tf.shape(input=x), tf.cast(low_fill, dtype)),
                      bins)

    if not extend_upper_interval:
      up_fill = np.nan if dtype.is_floating else tf.shape(input=edges)[0] - 1
      bins = tf.where(x > tf.expand_dims(edges[-1], 0),
                      tf.fill(tf.shape(input=x), tf.cast(up_fill, dtype)), bins)

    if flattening_x:
      bins = tf.reshape(bins, x_orig_shape)

    return bins


Instruction:
Can you implement a function in Python 3 that
returns a histogram of the number of events in a given axis.

Response:
def histogram(x,
              edges,
              axis=None,
              extend_lower_interval=False,
              extend_upper_interval=False,
              dtype=None,
              name=None):
  """Count how often `x` falls in intervals defined by `edges`.

  Given `edges = [c0, ..., cK]`, defining intervals
  `I0 = [c0, c1)`, `I1 = [c1, c2)`, ..., `I_{K-1} = [c_{K-1}, cK]`,
  This function counts how often `x` falls into each interval.

  Values of `x` outside of the intervals cause errors.  Consider using
  `extend_lower_interval`, `extend_upper_interval` to deal with this.

  Args:
    x:  Numeric `N-D` `Tensor` with `N > 0`.  If `axis` is not
      `None`, must have statically known number of dimensions. The
      `axis` kwarg determines which dimensions index iid samples.
      Other dimensions of `x` index "events" for which we will compute different
      histograms.
    edges:  `Tensor` of same `dtype` as `x`.  The first dimension indexes edges
      of intervals.  Must either be `1-D` or have `edges.shape[1:]` the same
      as the dimensions of `x` excluding `axis`.
      If `rank(edges) > 1`, `edges[k]` designates a shape `edges.shape[1:]`
      `Tensor` of interval edges for the corresponding dimensions of `x`.
    axis:  Optional `0-D` or `1-D` integer `Tensor` with constant
      values. The axis in `x` that index iid samples.
      `Default value:` `None` (treat every dimension as sample dimension).
    extend_lower_interval:  Python `bool`.  If `True`, extend the lowest
      interval `I0` to `(-inf, c1]`.
    extend_upper_interval:  Python `bool`.  If `True`, extend the upper
      interval `I_{K-1}` to `[c_{K-1}, +inf)`.
    dtype: The output type (`int32` or `int64`). `Default value:` `x.dtype`.
    name:  A Python string name to prepend to created ops.
      `Default value:` 'histogram'

  Returns:
    counts: `Tensor` of type `dtype` and, with
      `~axis = [i for i in range(arr.ndim) if i not in axis]`,
      `counts.shape = [edges.shape[0]] + x.shape[~axis]`.
      With `I` a multi-index into `~axis`, `counts[k][I]` is the number of times
      event(s) fell into the `kth` interval of `edges`.

  #### Examples

  ```python
  # x.shape = [1000, 2]
  # x[:, 0] ~ Uniform(0, 1), x[:, 1] ~ Uniform(1, 2).
  x = tf.stack([tf.random_uniform([1000]), 1 + tf.random_uniform([1000])],
               axis=-1)

  # edges ==> bins [0, 0.5), [0.5, 1.0), [1.0, 1.5), [1.5, 2.0].
  edges = [0., 0.5, 1.0, 1.5, 2.0]

  tfp.stats.histogram(x, edges)
  ==> approximately [500, 500, 500, 500]

  tfp.stats.histogram(x, edges, axis=0)
  ==> approximately [[500, 500, 0, 0], [0, 0, 500, 500]]
  ```

  """
  with tf.compat.v1.name_scope(name, 'histogram', values=[x, edges, axis]):

    # Tensor conversions.
    in_dtype = dtype_util.common_dtype([x, edges], preferred_dtype=tf.float32)

    x = tf.convert_to_tensor(value=x, name='x', dtype=in_dtype)
    edges = tf.convert_to_tensor(value=edges, name='edges', dtype=in_dtype)

    # Move dims in axis to the left end as one flattened dim.
    # After this, x.shape = [n_samples] + E.
    if axis is None:
      x = tf.reshape(x, shape=[-1])
    else:
      x_ndims = _get_static_ndims(
          x, expect_static=True, expect_ndims_at_least=1)
      axis = _make_static_axis_non_negative_list(axis, x_ndims)
      if not axis:
        raise ValueError('`axis` cannot be empty.  Found: {}'.format(axis))
      x = _move_dims_to_flat_end(x, axis, x_ndims, right_end=False)

    # bins.shape = x.shape = [n_samples] + E,
    # and bins[i] is a shape E Tensor of the bins that sample `i` fell into.
    # E is the "event shape", which is [] if axis is None.
    bins = find_bins(
        x,
        edges=edges,
        # If not extending intervals, then values outside the edges will return
        # -1, which gives an error when fed to bincount.
        extend_lower_interval=extend_lower_interval,
        extend_upper_interval=extend_upper_interval,
        dtype=tf.int32)

    # TODO(b/124015136) Use standard tf.math.bincount once it supports `axis`.
    counts = count_integers(
        bins,
        # Ensure we get correct output, even if x did not fall into every bin
        minlength=tf.shape(input=edges)[0] - 1,
        maxlength=tf.shape(input=edges)[0] - 1,
        axis=0,
        dtype=dtype or in_dtype)
    n_edges = tf.compat.dimension_value(edges.shape[0])
    if n_edges is not None:
      counts.set_shape(
          tf.TensorShape([n_edges - 1]).concatenate(counts.shape[1:]))
    return counts


Instruction:
Make a summary of the following Python 3 code
def percentile(x,
               q,
               axis=None,
               interpolation=None,
               keep_dims=False,
               validate_args=False,
               preserve_gradients=True,
               name=None):
  """Compute the `q`-th percentile(s) of `x`.

  Given a vector `x`, the `q`-th percentile of `x` is the value `q / 100` of the
  way from the minimum to the maximum in a sorted copy of `x`.

  The values and distances of the two nearest neighbors as well as the
  `interpolation` parameter will determine the percentile if the normalized
  ranking does not match the location of `q` exactly.

  This function is the same as the median if `q = 50`, the same as the minimum
  if `q = 0` and the same as the maximum if `q = 100`.

  Multiple percentiles can be computed at once by using `1-D` vector `q`.
  Dimension zero of the returned `Tensor` will index the different percentiles.

  Compare to `numpy.percentile`.

  Args:
    x:  Numeric `N-D` `Tensor` with `N > 0`.  If `axis` is not `None`,
      `x` must have statically known number of dimensions.
    q:  Scalar or vector `Tensor` with values in `[0, 100]`. The percentile(s).
    axis:  Optional `0-D` or `1-D` integer `Tensor` with constant values. The
      axis that index independent samples over which to return the desired
      percentile.  If `None` (the default), treat every dimension as a sample
      dimension, returning a scalar.
    interpolation : {'nearest', 'linear', 'lower', 'higher', 'midpoint'}.
      Default value: 'nearest'.  This specifies the interpolation method to
      use when the desired quantile lies between two data points `i < j`:
        * linear: i + (j - i) * fraction, where fraction is the fractional part
          of the index surrounded by i and j.
        * lower: `i`.
        * higher: `j`.
        * nearest: `i` or `j`, whichever is nearest.
        * midpoint: (i + j) / 2.
      `linear` and `midpoint` interpolation do not work with integer dtypes.
    keep_dims:  Python `bool`. If `True`, the last dimension is kept with size 1
      If `False`, the last dimension is removed from the output shape.
    validate_args:  Whether to add runtime checks of argument validity. If
      False, and arguments are incorrect, correct behavior is not guaranteed.
    preserve_gradients:  Python `bool`.  If `True`, ensure that gradient w.r.t
      the percentile `q` is preserved in the case of linear interpolation.
      If `False`, the gradient will be (incorrectly) zero when `q` corresponds
      to a point in `x`.
    name:  A Python string name to give this `Op`.  Default is 'percentile'

  Returns:
    A `(rank(q) + N - len(axis))` dimensional `Tensor` of same dtype as `x`, or,
      if `axis` is `None`, a `rank(q)` `Tensor`.  The first `rank(q)` dimensions
      index quantiles for different values of `q`.

  Raises:
    ValueError:  If argument 'interpolation' is not an allowed type.
    ValueError:  If interpolation type not compatible with `dtype`.

  #### Examples

  ```python
  # Get 30th percentile with default ('nearest') interpolation.
  x = [1., 2., 3., 4.]
  tfp.stats.percentile(x, q=30.)
  ==> 2.0

  # Get 30th percentile with 'linear' interpolation.
  x = [1., 2., 3., 4.]
  tfp.stats.percentile(x, q=30., interpolation='linear')
  ==> 1.9

  # Get 30th and 70th percentiles with 'lower' interpolation
  x = [1., 2., 3., 4.]
  tfp.stats.percentile(x, q=[30., 70.], interpolation='lower')
  ==> [1., 3.]

  # Get 100th percentile (maximum).  By default, this is computed over every dim
  x = [[1., 2.]
       [3., 4.]]
  tfp.stats.percentile(x, q=100.)
  ==> 4.

  # Treat the leading dim as indexing samples, and find the 100th quantile (max)
  # over all such samples.
  x = [[1., 2.]
       [3., 4.]]
  tfp.stats.percentile(x, q=100., axis=[0])
  ==> [3., 4.]
  ```

  """
  name = name or 'percentile'
  allowed_interpolations = {'linear', 'lower', 'higher', 'nearest', 'midpoint'}

  if interpolation is None:
    interpolation = 'nearest'
  else:
    if interpolation not in allowed_interpolations:
      raise ValueError('Argument `interpolation` must be in %s.  Found %s' %
                       (allowed_interpolations, interpolation))

  with tf.compat.v1.name_scope(name, values=[x, q]):
    x = tf.convert_to_tensor(value=x, name='x')

    if interpolation in {'linear', 'midpoint'} and x.dtype.is_integer:
      raise TypeError('{} interpolation not allowed with dtype {}'.format(
          interpolation, x.dtype))

    # Double is needed here and below, else we get the wrong index if the array
    # is huge along axis.
    q = tf.cast(q, tf.float64)
    _get_static_ndims(q, expect_ndims_no_more_than=1)

    if validate_args:
      q = distribution_util.with_dependencies([
          tf.compat.v1.assert_rank_in(q, [0, 1]),
          tf.compat.v1.assert_greater_equal(q, tf.cast(0., tf.float64)),
          tf.compat.v1.assert_less_equal(q, tf.cast(100., tf.float64))
      ], q)

    # Move `axis` dims of `x` to the rightmost, call it `y`.
    if axis is None:
      y = tf.reshape(x, [-1])
    else:
      x_ndims = _get_static_ndims(
          x, expect_static=True, expect_ndims_at_least=1)
      axis = _make_static_axis_non_negative_list(axis, x_ndims)
      y = _move_dims_to_flat_end(x, axis, x_ndims, right_end=True)

    frac_at_q_or_above = 1. - q / 100.

    # Sort everything, not just the top 'k' entries, which allows multiple calls
    # to sort only once (under the hood) and use CSE.
    sorted_y = _sort_tensor(y)

    d = tf.cast(tf.shape(input=y)[-1], tf.float64)

    def _get_indices(interp_type):
      """Get values of y at the indices implied by interp_type."""
      # Note `lower` <--> ceiling.  Confusing, huh?  Due to the fact that
      # _sort_tensor sorts highest to lowest, tf.ceil corresponds to the higher
      # index, but the lower value of y!
      if interp_type == 'lower':
        indices = tf.math.ceil((d - 1) * frac_at_q_or_above)
      elif interp_type == 'higher':
        indices = tf.floor((d - 1) * frac_at_q_or_above)
      elif interp_type == 'nearest':
        indices = tf.round((d - 1) * frac_at_q_or_above)
      # d - 1 will be distinct from d in int32, but not necessarily double.
      # So clip to avoid out of bounds errors.
      return tf.clip_by_value(
          tf.cast(indices, tf.int32), 0,
          tf.shape(input=y)[-1] - 1)

    if interpolation in ['nearest', 'lower', 'higher']:
      gathered_y = tf.gather(sorted_y, _get_indices(interpolation), axis=-1)
    elif interpolation == 'midpoint':
      gathered_y = 0.5 * (
          tf.gather(sorted_y, _get_indices('lower'), axis=-1) +
          tf.gather(sorted_y, _get_indices('higher'), axis=-1))
    elif interpolation == 'linear':
      # Copy-paste of docstring on interpolation:
      # linear: i + (j - i) * fraction, where fraction is the fractional part
      # of the index surrounded by i and j.
      larger_y_idx = _get_indices('lower')
      exact_idx = (d - 1) * frac_at_q_or_above
      if preserve_gradients:
        # If q corresponds to a point in x, we will initially have
        # larger_y_idx == smaller_y_idx.
        # This results in the gradient w.r.t. fraction being zero (recall `q`
        # enters only through `fraction`...and see that things cancel).
        # The fix is to ensure that smaller_y_idx and larger_y_idx are always
        # separated by exactly 1.
        smaller_y_idx = tf.maximum(larger_y_idx - 1, 0)
        larger_y_idx = tf.minimum(smaller_y_idx + 1, tf.shape(input=y)[-1] - 1)
        fraction = tf.cast(larger_y_idx, tf.float64) - exact_idx
      else:
        smaller_y_idx = _get_indices('higher')
        fraction = tf.math.ceil((d - 1) * frac_at_q_or_above) - exact_idx

      fraction = tf.cast(fraction, y.dtype)
      gathered_y = (
          tf.gather(sorted_y, larger_y_idx, axis=-1) * (1 - fraction) +
          tf.gather(sorted_y, smaller_y_idx, axis=-1) * fraction)

    # Propagate NaNs
    if x.dtype in (tf.bfloat16, tf.float16, tf.float32, tf.float64):
      # Apparently tf.is_nan doesn't like other dtypes
      nan_batch_members = tf.reduce_any(
          input_tensor=tf.math.is_nan(x), axis=axis)
      right_rank_matched_shape = tf.pad(
          tensor=tf.shape(input=nan_batch_members),
          paddings=[[0, tf.rank(input=q)]],
          constant_values=1)
      nan_batch_members = tf.reshape(
          nan_batch_members, shape=right_rank_matched_shape)
      shape_gathered_y = tf.shape(input=gathered_y)
      nan = np.array(np.nan, gathered_y.dtype.as_numpy_dtype)
      gathered_y = tf.where(
          tf.broadcast_to(nan_batch_members, shape_gathered_y),
          tf.fill(shape_gathered_y, nan),
          gathered_y)

    # Expand dimensions if requested
    if keep_dims:
      if axis is None:
        ones_vec = tf.ones(
            shape=[_get_best_effort_ndims(x) + _get_best_effort_ndims(q)],
            dtype=tf.int32)
        gathered_y *= tf.ones(ones_vec, dtype=x.dtype)
      else:
        gathered_y = _insert_back_keep_dims(gathered_y, axis)

    # If q is a scalar, then result has the right shape.
    # If q is a vector, then result has trailing dim of shape q.shape, which
    # needs to be rotated to dim 0.
    return distribution_util.rotate_transpose(gathered_y, tf.rank(q))

Response: Compute the percentile of x along a given axis.


Instruction:
Can you create a Python 3 function that
computes the quantiles of a distribution along a specified axis.

Response:
def quantiles(x,
              num_quantiles,
              axis=None,
              interpolation=None,
              keep_dims=False,
              validate_args=False,
              name=None):
  """Compute quantiles of `x` along `axis`.

  The quantiles of a distribution are cut points dividing the range into
  intervals with equal probabilities.

  Given a vector `x` of samples, this function estimates the cut points by
  returning `num_quantiles + 1` cut points, `(c0, ..., cn)`, such that, roughly
  speaking, equal number of sample points lie in the `num_quantiles` intervals
  `[c0, c1), [c1, c2), ..., [c_{n-1}, cn]`.  That is,

  * About `1 / n` fraction of the data lies in `[c_{k-1}, c_k)`, `k = 1, ..., n`
  * About `k / n` fraction of the data lies below `c_k`.
  * `c0` is the sample minimum and `cn` is the maximum.

  The exact number of data points in each interval depends on the size of
  `x` (e.g. whether the size is divisible by `n`) and the `interpolation` kwarg.

  Args:
    x:  Numeric `N-D` `Tensor` with `N > 0`.  If `axis` is not `None`,
      `x` must have statically known number of dimensions.
    num_quantiles:  Scalar `integer` `Tensor`.  The number of intervals the
      returned `num_quantiles + 1` cut points divide the range into.
    axis:  Optional `0-D` or `1-D` integer `Tensor` with constant values. The
      axis that index independent samples over which to return the desired
      percentile.  If `None` (the default), treat every dimension as a sample
      dimension, returning a scalar.
    interpolation : {'nearest', 'linear', 'lower', 'higher', 'midpoint'}.
      Default value: 'nearest'.  This specifies the interpolation method to
      use when the fractions `k / n` lie between two data points `i < j`:
        * linear: i + (j - i) * fraction, where fraction is the fractional part
          of the index surrounded by i and j.
        * lower: `i`.
        * higher: `j`.
        * nearest: `i` or `j`, whichever is nearest.
        * midpoint: (i + j) / 2. `linear` and `midpoint` interpolation do not
          work with integer dtypes.
    keep_dims:  Python `bool`. If `True`, the last dimension is kept with size 1
      If `False`, the last dimension is removed from the output shape.
    validate_args:  Whether to add runtime checks of argument validity. If
      False, and arguments are incorrect, correct behavior is not guaranteed.
    name:  A Python string name to give this `Op`.  Default is 'percentile'

  Returns:
    cut_points:  A `rank(x) + 1 - len(axis)` dimensional `Tensor` with same
    `dtype` as `x` and shape `[num_quantiles + 1, ...]` where the trailing shape
    is that of `x` without the dimensions in `axis` (unless `keep_dims is True`)

  Raises:
    ValueError:  If argument 'interpolation' is not an allowed type.
    ValueError:  If interpolation type not compatible with `dtype`.

  #### Examples

  ```python
  # Get quartiles of x with various interpolation choices.
  x = [0.,  1.,   2.,   3.,   4.,   5.,   6.,   7.,   8.,   9.,  10.]

  tfp.stats.quantiles(x, num_quantiles=4, interpolation='nearest')
  ==> [  0.,   2.,   5.,   8.,  10.]

  tfp.stats.quantiles(x, num_quantiles=4, interpolation='linear')
  ==> [  0. ,   2.5,   5. ,   7.5,  10. ]

  tfp.stats.quantiles(x, num_quantiles=4, interpolation='lower')
  ==> [  0.,   2.,   5.,   7.,  10.]

  # Get deciles of columns of an R x C data set.
  data = load_my_columnar_data(...)
  tfp.stats.quantiles(data, num_quantiles=10)
  ==> Shape [11, C] Tensor
  ```

  """
  with tf.compat.v1.name_scope(
      name, 'quantiles', values=[x, num_quantiles, axis]):
    x = tf.convert_to_tensor(value=x, name='x')
    return percentile(
        x,
        q=tf.linspace(
            # percentile casts q to float64 before using it...so may as well use
            # float64 here. Note that  using x.dtype won't work with linspace
            # if x is integral type (which is anothe motivation for hard-coding
            # float64).
            tf.convert_to_tensor(value=0, dtype=tf.float64),
            tf.convert_to_tensor(value=100, dtype=tf.float64),
            num=num_quantiles + 1),
        axis=axis,
        interpolation=interpolation,
        keep_dims=keep_dims,
        validate_args=validate_args,
        preserve_gradients=False)


Instruction:
Can you create a Python 3 function that
gets the number of dimensions and assert that some expectations are met.

Response:
def _get_static_ndims(x,
                      expect_static=False,
                      expect_ndims=None,
                      expect_ndims_no_more_than=None,
                      expect_ndims_at_least=None):
  """Get static number of dimensions and assert that some expectations are met.

  This function returns the number of dimensions 'ndims' of x, as a Python int.

  The optional expect arguments are used to check the ndims of x, but this is
  only done if the static ndims of x is not None.

  Args:
    x:  A Tensor.
    expect_static:  Expect `x` to have statically defined `ndims`.
    expect_ndims:  Optional Python integer.  If provided, assert that x has
      number of dimensions equal to this.
    expect_ndims_no_more_than:  Optional Python integer.  If provided, assert
      that x has no more than this many dimensions.
    expect_ndims_at_least:  Optional Python integer.  If provided, assert that x
      has at least this many dimensions.

  Returns:
    ndims:  A Python integer.

  Raises:
    ValueError:  If any of the expectations above are violated.
  """
  ndims = x.shape.ndims
  if ndims is None:
    shape_const = tf.get_static_value(tf.shape(input=x))
    if shape_const is not None:
      ndims = shape_const.ndim

  if ndims is None:
    if expect_static:
      raise ValueError(
          'Expected argument `x` to have statically defined `ndims`.  Found: ' %
          x)
    return

  if expect_ndims is not None:
    ndims_message = ('Expected argument `x` to have ndims %s.  Found tensor %s'
                     % (expect_ndims, x))
    if ndims != expect_ndims:
      raise ValueError(ndims_message)

  if expect_ndims_at_least is not None:
    ndims_at_least_message = (
        'Expected argument `x` to have ndims >= %d.  Found tensor %s' %
        (expect_ndims_at_least, x))
    if ndims < expect_ndims_at_least:
      raise ValueError(ndims_at_least_message)

  if expect_ndims_no_more_than is not None:
    ndims_no_more_than_message = (
        'Expected argument `x` to have ndims <= %d.  Found tensor %s' %
        (expect_ndims_no_more_than, x))
    if ndims > expect_ndims_no_more_than:
      raise ValueError(ndims_no_more_than_message)

  return ndims


Instruction:
Write a Python 3 script to
get ndims if possible.

Response:
def _get_best_effort_ndims(x,
                           expect_ndims=None,
                           expect_ndims_at_least=None,
                           expect_ndims_no_more_than=None):
  """Get static ndims if possible.  Fallback on `tf.rank(x)`."""
  ndims_static = _get_static_ndims(
      x,
      expect_ndims=expect_ndims,
      expect_ndims_at_least=expect_ndims_at_least,
      expect_ndims_no_more_than=expect_ndims_no_more_than)
  if ndims_static is not None:
    return ndims_static
  return tf.rank(x)


Instruction:
Create a Python 3 function to
insert the dims in axis back as singletons after being removed.

Response:
def _insert_back_keep_dims(x, axis):
  """Insert the dims in `axis` back as singletons after being removed.

  Args:
    x:  `Tensor`.
    axis:  Python list of integers.

  Returns:
    `Tensor` with same values as `x`, but additional singleton dimensions.
  """
  for i in sorted(axis):
    x = tf.expand_dims(x, axis=i)
  return x


Instruction:
How would you implement a function in Python 3 that
converts possibly negative indexed axis to non - negative list of ints.

Response:
def _make_static_axis_non_negative_list(axis, ndims):
  """Convert possibly negatively indexed axis to non-negative list of ints.

  Args:
    axis:  Integer Tensor.
    ndims:  Number of dimensions into which axis indexes.

  Returns:
    A list of non-negative Python integers.

  Raises:
    ValueError: If `axis` is not statically defined.
  """
  axis = distribution_util.make_non_negative_axis(axis, ndims)

  axis_const = tf.get_static_value(axis)
  if axis_const is None:
    raise ValueError(
        'Expected argument `axis` to be statically available.  Found: %s' %
        axis)

  # Make at least 1-D.
  axis = axis_const + np.zeros([1], dtype=axis_const.dtype)

  return list(int(dim) for dim in axis)


Instruction:
Write a Python 3 script to
move dims corresponding to axis in x to the end then flatten.

Response:
def _move_dims_to_flat_end(x, axis, x_ndims, right_end=True):
  """Move dims corresponding to `axis` in `x` to the end, then flatten.

  Args:
    x: `Tensor` with shape `[B0,B1,...,Bb]`.
    axis:  Python list of indices into dimensions of `x`.
    x_ndims:  Python integer holding number of dimensions in `x`.
    right_end:  Python bool.  Whether to move dims to the right end (else left).

  Returns:
    `Tensor` with value from `x` and dims in `axis` moved to end into one single
      dimension.
  """

  if not axis:
    return x

  # Suppose x.shape = [a, b, c, d]
  # Suppose axis = [1, 3]

  # other_dims = [0, 2] in example above.
  other_dims = sorted(set(range(x_ndims)).difference(axis))
  # x_permed.shape = [a, c, b, d]
  perm = other_dims + list(axis) if right_end else list(axis) + other_dims
  x_permed = tf.transpose(a=x, perm=perm)

  if x.shape.is_fully_defined():
    x_shape = x.shape.as_list()
    # other_shape = [a, c], end_shape = [b * d]
    other_shape = [x_shape[i] for i in other_dims]
    end_shape = [np.prod([x_shape[i] for i in axis])]
    full_shape = (
        other_shape + end_shape if right_end else end_shape + other_shape)
  else:
    other_shape = tf.gather(tf.shape(input=x), other_dims)
    full_shape = tf.concat(
        [other_shape, [-1]] if right_end else [[-1], other_shape], axis=0)
  return tf.reshape(x_permed, shape=full_shape)


Instruction:
Write a Python 3 function for
using top_k to sort a Tensor along the last dimension.

Response:
def _sort_tensor(tensor):
  """Use `top_k` to sort a `Tensor` along the last dimension."""
  sorted_, _ = tf.nn.top_k(tensor, k=tf.shape(input=tensor)[-1])
  sorted_.set_shape(tensor.shape)
  return sorted_


Instruction:
Write a Python 3 function that can
build a list of LinearGaussianStateSpaceModel objects for each component in the current state space.

Response:
def make_component_state_space_models(self,
                                        num_timesteps,
                                        param_vals,
                                        initial_step=0):
    """Build an ordered list of Distribution instances for component models.

    Args:
      num_timesteps: Python `int` number of timesteps to model.
      param_vals: a list of `Tensor` parameter values in order corresponding to
        `self.parameters`, or a dict mapping from parameter names to values.
      initial_step: optional `int` specifying the initial timestep to model.
        This is relevant when the model contains time-varying components,
        e.g., holidays or seasonality.

    Returns:
      component_ssms: a Python list of `LinearGaussianStateSpaceModel`
        Distribution objects, in order corresponding to `self.components`.
    """

    with tf.compat.v1.name_scope('make_component_state_space_models'):

      # List the model parameters in canonical order
      param_map = self._canonicalize_param_vals_as_map(param_vals)
      param_vals_list = [param_map[p.name] for p in self.parameters]

      # Build SSMs for each component model. We process the components in
      # canonical order, extracting the parameters for each component from the
      # (ordered) list of parameters.
      remaining_param_vals = param_vals_list[1:]
      component_ssms = []
      for component in self.components:
        num_parameters = len(component.parameters)
        component_param_vals = remaining_param_vals[:num_parameters]
        remaining_param_vals = remaining_param_vals[num_parameters:]

        component_ssms.append(
            component.make_state_space_model(
                num_timesteps,
                param_vals=component_param_vals,
                initial_step=initial_step))

    return component_ssms


Instruction:
Here you have a function in Python 3, explain what it does
def kl_reverse(logu, self_normalized=False, name=None):

  with tf.compat.v1.name_scope(name, "kl_reverse", [logu]):
    return amari_alpha(logu, alpha=0., self_normalized=self_normalized)

Response: A function that computes the inverse Kullback - Leibler Csiszar - function at logu.


Instruction:
Can you tell what is the following Python 3 function doing
def jensen_shannon(logu, self_normalized=False, name=None):

  with tf.compat.v1.name_scope(name, "jensen_shannon", [logu]):
    logu = tf.convert_to_tensor(value=logu, name="logu")
    npdt = logu.dtype.as_numpy_dtype
    y = tf.nn.softplus(logu)
    if self_normalized:
      y -= np.log(2).astype(npdt)
    return tf.exp(logu) * logu - (1. + tf.exp(logu)) * y

Response: A function that computes the Jensen - Shannon Csiszar - function in log - space.


Instruction:
How would you explain what the following Python 3 function does
def triangular(logu, name=None):

  with tf.compat.v1.name_scope(name, "triangular", [logu]):
    logu = tf.convert_to_tensor(value=logu, name="logu")
    return pearson(logu) / (1. + tf.exp(logu))

Response: A function that evaluates a Triangular Csiszar - function at logu.


Instruction:
Can you generate the documentation for the following Python 3 function
def t_power(logu, t, self_normalized=False, name=None):
  with tf.compat.v1.name_scope(name, "t_power", [logu, t]):
    logu = tf.convert_to_tensor(value=logu, name="logu")
    t = tf.convert_to_tensor(value=t, dtype=logu.dtype.base_dtype, name="t")
    fu = tf.math.expm1(t * logu)
    if self_normalized:
      fu -= t * tf.math.expm1(logu)
    fu *= tf.where(tf.logical_and(0. < t, t < 1.),
                   -tf.ones_like(t),
                   tf.ones_like(t))
    return fu

Response: A function that evaluates the T - Power Csiszar - function at logu at t.


Instruction:
Can you generate a brief explanation for the following Python 3 code
def modified_gan(logu, self_normalized=False, name=None):

  with tf.compat.v1.name_scope(name, "chi_square", [logu]):
    logu = tf.convert_to_tensor(value=logu, name="logu")
    y = tf.nn.softplus(logu) - logu
    if self_normalized:
      y += 0.5 * tf.math.expm1(logu)
    return y

Response: The Modified - GAN Csiszar - function in log - space.


Instruction:
How would you explain what the following Python 3 function does
def dual_csiszar_function(logu, csiszar_function, name=None):

  with tf.compat.v1.name_scope(name, "dual_csiszar_function", [logu]):
    return tf.exp(logu) * csiszar_function(-logu)

Response: Calculates the dual of a Csiszar - function over log - space.


Instruction:
Can you tell what is the following Python 3 function doing
def symmetrized_csiszar_function(logu, csiszar_function, name=None):

  with tf.compat.v1.name_scope(name, "symmetrized_csiszar_function", [logu]):
    logu = tf.convert_to_tensor(value=logu, name="logu")
    return 0.5 * (csiszar_function(logu)
                  + dual_csiszar_function(logu, csiszar_function))

Response: This function symmetrizes a Csiszar - function in log - space.


Instruction:
How would you explain what the following Python 3 function does
def monte_carlo_csiszar_f_divergence(
    f,
    p_log_prob,
    q,
    num_draws,
    use_reparametrization=None,
    seed=None,
    name=None):
  """Monte-Carlo approximation of the Csiszar f-Divergence.

  A Csiszar-function is a member of,

  ```none
  F = { f:R_+ to R : f convex }.
  ```

  The Csiszar f-Divergence for Csiszar-function f is given by:

  ```none
  D_f[p(X), q(X)] := E_{q(X)}[ f( p(X) / q(X) ) ]
                  ~= m**-1 sum_j^m f( p(x_j) / q(x_j) ),
                             where x_j ~iid q(X)
  ```

  Tricks: Reparameterization and Score-Gradient

  When q is "reparameterized", i.e., a diffeomorphic transformation of a
  parameterless distribution (e.g.,
  `Normal(Y; m, s) <=> Y = sX + m, X ~ Normal(0,1)`), we can swap gradient and
  expectation, i.e.,
  `grad[Avg{ s_i : i=1...n }] = Avg{ grad[s_i] : i=1...n }` where `S_n=Avg{s_i}`
  and `s_i = f(x_i), x_i ~iid q(X)`.

  However, if q is not reparameterized, TensorFlow's gradient will be incorrect
  since the chain-rule stops at samples of unreparameterized distributions. In
  this circumstance using the Score-Gradient trick results in an unbiased
  gradient, i.e.,

  ```none
  grad[ E_q[f(X)] ]
  = grad[ int dx q(x) f(x) ]
  = int dx grad[ q(x) f(x) ]
  = int dx [ q'(x) f(x) + q(x) f'(x) ]
  = int dx q(x) [q'(x) / q(x) f(x) + f'(x) ]
  = int dx q(x) grad[ f(x) q(x) / stop_grad[q(x)] ]
  = E_q[ grad[ f(x) q(x) / stop_grad[q(x)] ] ]
  ```

  Unless `q.reparameterization_type != tfd.FULLY_REPARAMETERIZED` it is
  usually preferable to set `use_reparametrization = True`.

  Example Application:

  The Csiszar f-Divergence is a useful framework for variational inference.
  I.e., observe that,

  ```none
  f(p(x)) =  f( E_{q(Z | x)}[ p(x, Z) / q(Z | x) ] )
          <= E_{q(Z | x)}[ f( p(x, Z) / q(Z | x) ) ]
          := D_f[p(x, Z), q(Z | x)]
  ```

  The inequality follows from the fact that the "perspective" of `f`, i.e.,
  `(s, t) |-> t f(s / t))`, is convex in `(s, t)` when `s/t in domain(f)` and
  `t` is a real. Since the above framework includes the popular Evidence Lower
  BOund (ELBO) as a special case, i.e., `f(u) = -log(u)`, we call this framework
  "Evidence Divergence Bound Optimization" (EDBO).

  Args:
    f: Python `callable` representing a Csiszar-function in log-space, i.e.,
      takes `p_log_prob(q_samples) - q.log_prob(q_samples)`.
    p_log_prob: Python `callable` taking (a batch of) samples from `q` and
      returning the natural-log of the probability under distribution `p`.
      (In variational inference `p` is the joint distribution.)
    q: `tf.Distribution`-like instance; must implement:
      `reparameterization_type`, `sample(n, seed)`, and `log_prob(x)`.
      (In variational inference `q` is the approximate posterior distribution.)
    num_draws: Integer scalar number of draws used to approximate the
      f-Divergence expectation.
    use_reparametrization: Python `bool`. When `None` (the default),
      automatically set to:
      `q.reparameterization_type == tfd.FULLY_REPARAMETERIZED`.
      When `True` uses the standard Monte-Carlo average. When `False` uses the
      score-gradient trick. (See above for details.)  When `False`, consider
      using `csiszar_vimco`.
    seed: Python `int` seed for `q.sample`.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    monte_carlo_csiszar_f_divergence: `float`-like `Tensor` Monte Carlo
      approximation of the Csiszar f-Divergence.

  Raises:
    ValueError: if `q` is not a reparameterized distribution and
      `use_reparametrization = True`. A distribution `q` is said to be
      "reparameterized" when its samples are generated by transforming the
      samples of another distribution which does not depend on the
      parameterization of `q`. This property ensures the gradient (with respect
      to parameters) is valid.
    TypeError: if `p_log_prob` is not a Python `callable`.
  """
  reparameterization_types = tf.nest.flatten(q.reparameterization_type)
  with tf.compat.v1.name_scope(name, "monte_carlo_csiszar_f_divergence",
                               [num_draws]):
    if use_reparametrization is None:
      use_reparametrization = all(
          reparameterization_type == tfd.FULLY_REPARAMETERIZED
          for reparameterization_type in reparameterization_types)
    elif (use_reparametrization and
          any(reparameterization_type != tfd.FULLY_REPARAMETERIZED
              for reparameterization_type in reparameterization_types)):
      # TODO(jvdillon): Consider only raising an exception if the gradient is
      # requested.
      raise ValueError(
          "Distribution `q` must be reparameterized, i.e., a diffeomorphic "
          "transformation of a parameterless distribution. (Otherwise this "
          "function has a biased gradient.)")
    if not callable(p_log_prob):
      raise TypeError("`p_log_prob` must be a Python `callable` function.")
    return monte_carlo.expectation(
        f=lambda q_samples: f(p_log_prob(q_samples) - q.log_prob(q_samples)),
        samples=q.sample(num_draws, seed=seed),
        log_prob=q.log_prob,  # Only used if use_reparametrization=False.
        use_reparametrization=use_reparametrization)

Response: Monte - Carlo approximation of the Csiszar - function f - Divergence.


Instruction:
Here you have a function in Python 3, explain what it does
def csiszar_vimco(f,
                  p_log_prob,
                  q,
                  num_draws,
                  num_batch_draws=1,
                  seed=None,
                  name=None):
  """Use VIMCO to lower the variance of gradient[csiszar_function(Avg(logu))].

  This function generalizes VIMCO [(Mnih and Rezende, 2016)][1] to Csiszar
  f-Divergences.

  Note: if `q.reparameterization_type = tfd.FULLY_REPARAMETERIZED`,
  consider using `monte_carlo_csiszar_f_divergence`.

  The VIMCO loss is:

  ```none
  vimco = f(Avg{logu[i] : i=0,...,m-1})
  where,
    logu[i] = log( p(x, h[i]) / q(h[i] | x) )
    h[i] iid~ q(H | x)
  ```

  Interestingly, the VIMCO gradient is not the naive gradient of `vimco`.
  Rather, it is characterized by:

  ```none
  grad[vimco] - variance_reducing_term
  where,
    variance_reducing_term = Sum{ grad[log q(h[i] | x)] *
                                    (vimco - f(log Avg{h[j;i] : j=0,...,m-1}))
                                 : i=0, ..., m-1 }
    h[j;i] = { u[j]                             j!=i
             { GeometricAverage{ u[k] : k!=i}   j==i
  ```

  (We omitted `stop_gradient` for brevity. See implementation for more details.)

  The `Avg{h[j;i] : j}` term is a kind of "swap-out average" where the `i`-th
  element has been replaced by the leave-`i`-out Geometric-average.

  This implementation prefers numerical precision over efficiency, i.e.,
  `O(num_draws * num_batch_draws * prod(batch_shape) * prod(event_shape))`.
  (The constant may be fairly large, perhaps around 12.)

  Args:
    f: Python `callable` representing a Csiszar-function in log-space.
    p_log_prob: Python `callable` representing the natural-log of the
      probability under distribution `p`. (In variational inference `p` is the
      joint distribution.)
    q: `tf.Distribution`-like instance; must implement: `sample(n, seed)`, and
      `log_prob(x)`. (In variational inference `q` is the approximate posterior
      distribution.)
    num_draws: Integer scalar number of draws used to approximate the
      f-Divergence expectation.
    num_batch_draws: Integer scalar number of draws used to approximate the
      f-Divergence expectation.
    seed: Python `int` seed for `q.sample`.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    vimco: The Csiszar f-Divergence generalized VIMCO objective.

  Raises:
    ValueError: if `num_draws < 2`.

  #### References

  [1]: Andriy Mnih and Danilo Rezende. Variational Inference for Monte Carlo
       objectives. In _International Conference on Machine Learning_, 2016.
       https://arxiv.org/abs/1602.06725
  """
  with tf.compat.v1.name_scope(name, "csiszar_vimco",
                               [num_draws, num_batch_draws]):
    if num_draws < 2:
      raise ValueError("Must specify num_draws > 1.")
    stop = tf.stop_gradient  # For readability.
    x = stop(q.sample(sample_shape=[num_draws, num_batch_draws],
                      seed=seed))
    logqx = q.log_prob(x)
    logu = p_log_prob(x) - logqx
    f_log_avg_u, f_log_sooavg_u = [f(r) for r in csiszar_vimco_helper(logu)]
    dotprod = tf.reduce_sum(
        input_tensor=logqx * stop(f_log_avg_u - f_log_sooavg_u),
        axis=0)  # Sum over iid samples.
    # We now rewrite f_log_avg_u so that:
    #   `grad[f_log_avg_u] := grad[f_log_avg_u + dotprod]`.
    # To achieve this, we use a trick that
    #   `f(x) - stop(f(x)) == zeros_like(f(x))`
    # but its gradient is grad[f(x)].
    # Note that IEEE754 specifies that `x - x == 0.` and `x + 0. == x`, hence
    # this trick loses no precision. For more discussion regarding the relevant
    # portions of the IEEE754 standard, see the StackOverflow question,
    # "Is there a floating point value of x, for which x-x == 0 is false?"
    # http://stackoverflow.com/q/2686644
    f_log_avg_u += dotprod - stop(dotprod)  # Add zeros_like(dot_prod).
    return tf.reduce_mean(input_tensor=f_log_avg_u, axis=0)

Response: This function is used to lower the variance of gradient of vimco.