def Func(arg_0,
                   arg_1,
                   arg_2=None,
                   arg_3='target_log_prob'):
  """Helper which computes `fn_result` if needed."""
  arg_1 = (list(arg_1) if mcmc_util.is_list_like(arg_1)
                 else [arg_1])
  if arg_2 is None:
    arg_2 = arg_0(*arg_1)
  if not arg_2.dtype.is_floating:
    raise TypeError('`{}` must be a `Tensor` with `float` `dtype`.'.format(
        arg_3))
  return arg_2


def Func(arg_0, arg_1):
  """Pads the shape of x to the right to be of rank final_rank.

  Expands the dims of `x` to the right such that its rank is equal to
  final_rank. For example, if `x` is of shape [1, 5, 7, 2] and `final_rank` is
  7, we return padded_x, which is of shape [1, 5, 7, 2, 1, 1, 1].

  Args:
    x: The tensor whose shape is to be padded.
    final_rank: Scalar int32 `Tensor` or Python `int`. The desired rank of x.

  Returns:
    padded_x: A tensor of rank final_rank.
  """
  arg_2 = tf.concat(
      [tf.shape(input=arg_0),
       tf.ones(arg_1 - tf.rank(arg_0), dtype=tf.int32)],
      axis=0)
  arg_3 = None
  if arg_0.shape.is_fully_defined() and isinstance(arg_1, int):
    arg_3 = arg_0.shape.as_list()
    arg_4 = arg_1 - len(arg_3)
    arg_3.extend([1] * arg_4)

  arg_5 = tf.reshape(arg_0, arg_3 or arg_2)
  return arg_5


def Func(arg_0, arg_1, arg_2):
    """Runs one iteration of Slice Sampler.

    Args:
      current_state: `Tensor` or Python `list` of `Tensor`s representing the
        current state(s) of the Markov chain(s). The first `r` dimensions
        index independent chains,
        `r = tf.rank(target_log_prob_fn(*current_state))`.
      previous_kernel_results: `collections.namedtuple` containing `Tensor`s
        representing values from previous calls to this function (or from the
        `bootstrap_results` function.)

    Returns:
      next_state: Tensor or Python list of `Tensor`s representing the state(s)
        of the Markov chain(s) after taking exactly one step. Has same type and
        shape as `current_state`.
      kernel_results: `collections.namedtuple` of internal calculations used to
        advance the chain.

    Raises:
      ValueError: if there isn't one `step_size` or a list with same length as
        `current_state`.
      TypeError: if `not target_log_prob.dtype.is_floating`.
    """
    with tf.compat.v1.name_scope(
        name=mcmc_util.make_name(arg_0.name, 'slice', 'Func'),
        values=[
            arg_0.step_size, arg_0.max_doublings, arg_0._seed_stream,
            arg_1, arg_2.target_log_prob
        ]):
      with tf.compat.v1.name_scope('initialize'):
        [
            arg_3,
            arg_4,
            arg_5
        ] = _prepare_args(
            arg_0.target_log_prob_fn,
            arg_1,
            arg_0.step_size,
            arg_2.target_log_prob,
            maybe_expand=True)

        arg_6 = tf.convert_to_tensor(
            value=arg_0.max_doublings, dtype=tf.int32, name='max_doublings')

      arg_7 = distribution_util.prefer_static_rank(
          arg_5)

      [
          arg_8,
          arg_9,
          arg_10,
          arg_11,
          arg_12,
          arg_13
      ] = _sample_next(
          arg_0.target_log_prob_fn,
          arg_3,
          arg_4,
          arg_6,
          arg_5,
          arg_7,
          seed=arg_0._seed_stream()
      )

      def maybe_flatten(arg_14):
        return arg_14 if mcmc_util.is_list_like(arg_1) else arg_14[0]

      return [
          maybe_flatten(arg_8),
          SliceSamplerKernelResults(
              target_log_prob=arg_9,
              arg_10=arg_10,
              arg_11=arg_11,
              arg_12=arg_12,
              arg_13=arg_13
          ),
      ]


def Func(arg_0,
                                    arg_1,
                                    arg_2=(),
                                    arg_3=None,
                                    arg_4=None):
  """Build a loss function for variational inference in STS models.

  Variational inference searches for the distribution within some family of
  approximate posteriors that minimizes a divergence between the approximate
  posterior `q(z)` and true posterior `p(z|observed_time_series)`. By converting
  inference to optimization, it's generally much faster than sampling-based
  inference algorithms such as HMC. The tradeoff is that the approximating
  family rarely contains the true posterior, so it may miss important aspects of
  posterior structure (in particular, dependence between variables) and should
  not be blindly trusted. Results may vary; it's generally wise to compare to
  HMC to evaluate whether inference quality is sufficient for your task at hand.

  This method constructs a loss function for variational inference using the
  Kullback-Liebler divergence `KL[q(z) || p(z|observed_time_series)]`, with an
  approximating family given by independent Normal distributions transformed to
  the appropriate parameter space for each parameter. Minimizing this loss (the
  negative ELBO) maximizes a lower bound on the log model evidence `-log
  p(observed_time_series)`. This is equivalent to the 'mean-field' method
  implemented in [1]. and is a standard approach. The resulting posterior
  approximations are unimodal; they will tend to underestimate posterior
  uncertainty when the true posterior contains multiple modes (the `KL[q||p]`
  divergence encourages choosing a single mode) or dependence between variables.

  Args:
    model: An instance of `StructuralTimeSeries` representing a
      time-series model. This represents a joint distribution over
      time-series and their parameters with batch shape `[b1, ..., bN]`.
    observed_time_series: `float` `Tensor` of shape
      `concat([sample_shape, model.batch_shape, [num_timesteps, 1]]) where
      `sample_shape` corresponds to i.i.d. observations, and the trailing `[1]`
      dimension may (optionally) be omitted if `num_timesteps > 1`. May
      optionally be an instance of `tfp.sts.MaskedTimeSeries`, which includes
      a mask `Tensor` to specify timesteps with missing observations.
    init_batch_shape: Batch shape (Python `tuple`, `list`, or `int`) of initial
      states to optimize in parallel.
      Default value: `()`. (i.e., just run a single optimization).
    seed: Python integer to seed the random number generator.
    name: Python `str` name prefixed to ops created by this function.
      Default value: `None` (i.e., 'Func').

  Returns:
    variational_loss: `float` `Tensor` of shape
      `concat([init_batch_shape, model.batch_shape])`, encoding a stochastic
      estimate of an upper bound on the negative model evidence `-log p(y)`.
      Minimizing this loss performs variational inference; the gap between the
      variational bound and the true (generally unknown) model evidence
      corresponds to the divergence `KL[q||p]` between the approximate and true
      posterior.
    variational_distributions: `collections.OrderedDict` giving
      the approximate posterior for each model parameter. The keys are
      Python `str` parameter names in order, corresponding to
      `[param.name for param in model.parameters]`. The values are
      `tfd.Distribution` instances with batch shape
      `concat([init_batch_shape, model.batch_shape])`; these will typically be
      of the form `tfd.TransformedDistribution(tfd.Normal(...),
      bijector=param.bijector)`.

  #### Examples

  Assume we've built a structural time-series model:

  ```python
    day_of_week = tfp.sts.Seasonal(
        num_seasons=7,
        observed_time_series=observed_time_series,
        name='day_of_week')
    local_linear_trend = tfp.sts.LocalLinearTrend(
        observed_time_series=observed_time_series,
        name='local_linear_trend')
    model = tfp.sts.Sum(components=[day_of_week, local_linear_trend],
                        observed_time_series=observed_time_series)
  ```

  To run variational inference, we simply construct the loss and optimize
  it:

  ```python
    (variational_loss,
     variational_distributions) = tfp.sts.Func(
       model=model, observed_time_series=observed_time_series)

    train_op = tf.train.AdamOptimizer(0.1).minimize(variational_loss)
    with tf.Session() as sess:
      sess.run(tf.global_variables_initializer())

      for step in range(200):
        _, loss_ = sess.run((train_op, variational_loss))

        if step % 20 == 0:
          print("step {} loss {}".format(step, loss_))

      posterior_samples_ = sess.run({
        param_name: q.sample(50)
        for param_name, q in variational_distributions.items()})
  ```

  As a more complex example, we might try to avoid local optima by optimizing
  from multiple initializations in parallel, and selecting the result with the
  lowest loss:

  ```python
    (variational_loss,
     variational_distributions) = tfp.sts.Func(
       model=model, observed_time_series=observed_time_series,
       init_batch_shape=[10])

    train_op = tf.train.AdamOptimizer(0.1).minimize(variational_loss)
    with tf.Session() as sess:
      sess.run(tf.global_variables_initializer())

      for step in range(200):
        _, loss_ = sess.run((train_op, variational_loss))

        if step % 20 == 0:
          print("step {} losses {}".format(step, loss_))

      # Draw multiple samples to reduce Monte Carlo error in the optimized
      # variational bounds.
      avg_loss = np.mean(
        [sess.run(variational_loss) for _ in range(25)], axis=0)
      best_posterior_idx = np.argmin(avg_loss, axis=0).astype(np.int32)
  ```

  #### References

  [1]: Alp Kucukelbir, Dustin Tran, Rajesh Ranganath, Andrew Gelman, and
       David M. Blei. Automatic Differentiation Variational Inference. In
       _Journal of Machine Learning Research_, 2017.
       https://arxiv.org/abs/1603.00788

  """

  with tf.compat.v1.name_scope(
      arg_4, 'Func',
      values=[arg_1]) as arg_4:
    arg_3 = tfd.SeedStream(
        arg_3, salt='StructuralTimeSeries_Func')

    arg_5 = collections.OrderedDict()
    arg_6 = []
    for arg_7 in arg_0.parameters:
      def initial_loc_fn(arg_7):
        return sample_uniform_initial_state(
            arg_7, return_constrained=True,
            init_sample_shape=arg_2,
            arg_3=arg_3())
      arg_8 = _build_trainable_posterior(arg_7, initial_loc_fn=initial_loc_fn)
      arg_5[arg_7.name] = arg_8
      arg_6.append(arg_8.sample(arg_3=arg_3()))

    # Multiple initializations (similar to HMC chains) manifest as an extra
    # param batch dimension, so we need to add corresponding batch dimension(s)
    # to `observed_time_series`.
    arg_1 = sts_util.pad_batch_dimension_for_multiple_chains(
        arg_1, arg_0, chain_batch_shape=arg_2)

    # Construct the variational bound.
    arg_9 = arg_0.joint_log_prob(arg_1)
    arg_10 = arg_9(*arg_6)
    arg_11 = tf.reduce_sum(
        input_tensor=[
            -arg_8.log_prob(sample) for (arg_8, sample) in zip(
                arg_5.values(), arg_6)
        ],
        axis=0)
    arg_12 = -(arg_10 + arg_11)  # -ELBO

  return arg_12, arg_5


def Func(arg_0, arg_1=200, arg_2=None):
  """Run an optimizer within the graph to minimize a loss function."""
  arg_2 = tf.compat.v1.train.AdamOptimizer(
      0.1) if arg_2 is None else arg_2

  def train_loop_body(arg_3):
    arg_4 = arg_2.minimize(
        arg_0 if tf.executing_eagerly() else arg_0())
    return tf.tuple(tensors=[tf.add(arg_3, 1)], control_inputs=[arg_4])

  arg_5 = tf.compat.v1.while_loop(
      cond=lambda arg_3: arg_3 < arg_1,
      body=train_loop_body,
      loop_vars=[tf.constant(0)],
      return_same_structure=True)[0]  # Always return a single op.
  return arg_5


def Func(arg_0, arg_1):
  """Compute mean and variance, accounting for a mask.

  Args:
    time_series_tensor: float `Tensor` time series of shape
      `concat([batch_shape, [num_timesteps]])`.
    broadcast_mask: bool `Tensor` of the same shape as `time_series`.
  Returns:
    mean: float `Tensor` of shape `batch_shape`.
    variance: float `Tensor` of shape `batch_shape`.
  """
  arg_2 = tf.cast(
      tf.reduce_sum(input_tensor=tf.cast(~arg_1, tf.int32), axis=-1),
      arg_0.dtype)

  # Manually compute mean and variance, excluding masked entries.
  arg_3 = (tf.reduce_sum(input_tensor=tf.where(
      arg_1,
      tf.zeros_like(arg_0),
      arg_0), axis=-1) / arg_2)
  arg_4 = (tf.reduce_sum(input_tensor=tf.where(
      arg_1,
      tf.zeros_like(arg_0),
      (arg_0 - arg_3[..., tf.newaxis]) ** 2), axis=-1)
              / arg_2)
  return arg_3, arg_4


def Func(arg_0):
  """Get broadcast batch shape from distributions, statically if possible."""

  # Static case
  arg_1 = arg_0[0].batch_shape
  for arg_2 in arg_0:
    arg_1 = tf.broadcast_static_shape(arg_1,
                                            arg_2.batch_shape)
  if arg_1.is_fully_defined():
    return arg_1.as_list()

  # Fallback on dynamic.
  arg_1 = arg_0[0].batch_shape_tensor()
  for arg_2 in arg_0:
    arg_1 = tf.broadcast_dynamic_shape(arg_1,
                                             arg_2.batch_shape_tensor())

  return tf.convert_to_tensor(value=arg_1)


def Func(arg_0):
  """Combine MultivariateNormals into a factored joint distribution.

   Given a list of multivariate normal distributions
   `dist[i] = Normal(loc[i], scale[i])`, construct the joint
   distribution given by concatenating independent samples from these
   distributions. This is multivariate normal with mean vector given by the
   concatenation of the component mean vectors, and block-diagonal covariance
   matrix in which the blocks are the component covariances.

   Note that for computational efficiency, multivariate normals are represented
   by a 'scale' (factored covariance) linear operator rather than the full
   covariance matrix.

  Args:
    distributions: Python `iterable` of MultivariateNormal distribution
      instances (e.g., `tfd.MultivariateNormalDiag`,
      `tfd.MultivariateNormalTriL`, etc.). These must be broadcastable to a
      consistent batch shape, but may have different event shapes
      (i.e., defined over spaces of different dimension).

  Returns:
    joint_distribution: An instance of `tfd.MultivariateNormalLinearOperator`
      representing the joint distribution constructed by concatenating
      an independent sample from each input distributions.
  """

  arg_1 = [tensor for distribution in arg_0
                   for tensor in distribution._graph_parents]  # pylint: disable=protected-access
  with tf.compat.v1.name_scope('Func', values=arg_1):

    # We explicitly broadcast the `locs` so that we can concatenate them.
    # We don't have direct numerical access to the `scales`, which are arbitrary
    # linear operators, but `LinearOperatorBlockDiag` appears to do the right
    # thing without further intervention.
    arg_2 = tf.debugging.assert_same_float_dtype(arg_0)
    arg_3 = tf.ones(broadcast_batch_shape(arg_0),
                             arg_2=arg_2)[..., tf.newaxis]
    return MultivariateNormalLinearOperator(
        loc=tf.concat([arg_4.mean() * arg_3 for arg_4 in arg_0],
                      axis=-1),
        scale=tfl.LinearOperatorBlockDiag([arg_4.scale for arg_4 in arg_0],
                                          is_square=True))


def Func(arg_0):
  """Compute statistics of a provided time series, as heuristic initialization.

  Args:
    observed_time_series: `Tensor` representing a time series, or batch of time
       series, of shape either `batch_shape + [num_timesteps, 1]` or
       `batch_shape + [num_timesteps]` (allowed if `num_timesteps > 1`).

  Returns:
    observed_mean: `Tensor` of shape `batch_shape`, giving the empirical
      mean of each time series in the batch.
    observed_stddev: `Tensor` of shape `batch_shape`, giving the empirical
      standard deviation of each time series in the batch.
    observed_initial_centered: `Tensor of shape `batch_shape`, giving the
      initial value of each time series in the batch after centering
      (subtracting the mean).
  """

  with tf.compat.v1.name_scope(
      'Func', values=[arg_0]):

    [
        arg_0,
        arg_1
    ] = canonicalize_observed_time_series_with_mask(arg_0)

    arg_2 = arg_0[..., 0]
    if arg_1 is None:
      arg_3, arg_4 = tf.nn.moments(
          x=arg_2, axes=-1)
      arg_5 = arg_2[..., 0]
    else:
      arg_6 = tf.broadcast_to(tf.cast(arg_1, tf.bool),
                                       tf.shape(input=arg_2))
      arg_3, arg_4 = (
          missing_values_util.moments_of_masked_time_series(
              arg_2, arg_6=arg_6))
      try:
        arg_5 = (
            missing_values_util.initial_value_of_masked_time_series(
                arg_2, arg_6=arg_6))
      except NotImplementedError:
        tf.compat.v1.logging.warn(
            'Cannot compute initial values for a masked time series'
            'with dynamic shape; using the mean instead. This will'
            'affect heuristic priors and may change the results of'
            'inference.')
        arg_5 = arg_3

    arg_7 = tf.sqrt(arg_4)
    arg_8 = arg_5 - arg_3
    return arg_3, arg_7, arg_8


def Func(arg_0):
  """Ensures `observed_time_series_tensor` has a trailing dimension of size 1.

  The `tfd.LinearGaussianStateSpaceModel` Distribution has event shape of
  `[num_timesteps, observation_size]`, but canonical BSTS models
  are univariate, so their observation_size is always `1`. The extra trailing
  dimension gets annoying, so this method allows arguments with or without the
  extra dimension. There is no ambiguity except in the trivial special case
  where  `num_timesteps = 1`; this can be avoided by specifying any unit-length
  series in the explicit `[num_timesteps, 1]` style.

  Most users should not call this method directly, and instead call
  `canonicalize_observed_time_series_with_mask`, which handles converting
  to `Tensor` and specifying an optional missingness mask.

  Args:
    observed_time_series_tensor: `Tensor` of shape
      `batch_shape + [num_timesteps, 1]` or `batch_shape + [num_timesteps]`,
      where `num_timesteps > 1`.

  Returns:
    expanded_time_series: `Tensor` of shape `batch_shape + [num_timesteps, 1]`.
  """

  with tf.compat.v1.name_scope(
      'maybe_expand_trailing_dim', values=[arg_0]):
    if (arg_0.shape.ndims is not None and
        tf.compat.dimension_value(
            arg_0.shape[-1]) is not None):
      arg_1 = (
          arg_0
          if arg_0.shape[-1] == 1
          else arg_0[..., tf.newaxis])
    else:
      arg_1 = tf.cond(
          pred=tf.equal(tf.shape(input=arg_0)[-1], 1),
          true_fn=lambda: arg_0,
          false_fn=lambda: arg_0[..., tf.newaxis])
    return arg_1


def Func(arg_0, arg_1):
  """Construct a predictive normal distribution that mixes over posterior draws.

  Args:
    means: float `Tensor` of shape
      `[num_posterior_draws, ..., num_timesteps]`.
    variances: float `Tensor` of shape
      `[num_posterior_draws, ..., num_timesteps]`.

  Returns:
    mixture_dist: `tfd.MixtureSameFamily(tfd.Independent(tfd.Normal))` instance
      representing a uniform mixture over the posterior samples, with
      `batch_shape = ...` and `event_shape = [num_timesteps]`.

  """
  # The inputs `means`, `variances` have shape
  #   `concat([
  #      [num_posterior_draws],
  #      sample_shape,
  #      batch_shape,
  #      [num_timesteps]])`
  # Because MixtureSameFamily mixes over the rightmost batch dimension,
  # we need to move the `num_posterior_draws` dimension to be rightmost
  # in the batch shape. This requires use of `Independent` (to preserve
  # `num_timesteps` as part of the event shape) and `move_dimension`.
  # TODO(b/120245392): enhance `MixtureSameFamily` to reduce along an
  # arbitrary axis, and eliminate `move_dimension` calls here.

  with tf.compat.v1.name_scope(
      'Func', values=[arg_0, arg_1]):
    arg_2 = dist_util.prefer_static_value(
        tf.shape(input=arg_0))[0]

    arg_3 = tfd.Independent(
        distribution=tfd.Normal(
            loc=dist_util.move_dimension(arg_0, 0, -2),
            scale=tf.sqrt(dist_util.move_dimension(arg_1, 0, -2))),
        reinterpreted_batch_ndims=1)

    return tfd.MixtureSameFamily(
        mixture_distribution=tfd.Categorical(
            logits=tf.zeros([arg_2],
                            dtype=arg_3.dtype)),
        components_distribution=arg_3)


def Func(arg_0, arg_1, arg_2):
  """Uses arg names to resolve distribution names."""
  if arg_1 is None:
    arg_1 = []
  else:
    arg_1 = arg_1.copy()
  arg_3 = len(arg_0)
  arg_1.extend([None]*(arg_3 - len(arg_1)))
  for arg_4, arg_5 in enumerate(reversed(arg_0)):
    if not arg_5:
      continue  # There's no args to analyze.
    arg_6 = arg_3 - arg_4 - 1
    for arg_7, arg_8 in enumerate(arg_5):
      arg_1[arg_6 - arg_7 - 1] = arg_8
  arg_7 = 0
  for arg_4 in range(len(arg_1)):
    arg_6 = arg_3 - arg_4 - 1
    if arg_1[arg_6] is None:
      arg_1[arg_6] = arg_2 if arg_7 == 0 else arg_2 + str(arg_7)
      arg_7 += 1
  return tuple(arg_1)


def Func(arg_0, arg_1, arg_2=None):
  """Calculate the KL divergence between two `JointDistributionSequential`s.

  Args:
    d0: instance of a `JointDistributionSequential` object.
    d1: instance of a `JointDistributionSequential` object.
    name: (optional) Name to use for created operations.
      Default value: `"kl_joint_joint"`.

  Returns:
    kl_joint_joint: `Tensor` The sum of KL divergences between elemental
      distributions of two joint distributions.

  Raises:
    ValueError: when joint distributions have a different number of elemental
      distributions.
    ValueError: when either joint distribution has a distribution with dynamic
      dependency, i.e., when either joint distribution is not a collection of
      independent distributions.
  """
  if len(arg_0._dist_fn_wrapped) != len(arg_1._dist_fn_wrapped):  # pylint: disable=protected-access
    raise ValueError(
        'Can only compute KL divergence between when each has the'
        'same number of component distributions.')
  if (not all(arg_3 is None for arg_3 in arg_0._dist_fn_args) or  # pylint: disable=protected-access
      not all(arg_3 is None for arg_3 in arg_1._dist_fn_args)):  # pylint: disable=protected-access
    raise ValueError(
        'Can only compute KL divergence when all distributions are '
        'independent.')
  with tf.name_scope(arg_2 or 'kl_jointseq_jointseq'):
    return sum(kullback_leibler.kl_divergence(arg_4(), arg_5())
               for arg_4, arg_5 in zip(arg_0._dist_fn_wrapped, arg_1._dist_fn_wrapped))


def Func(arg_0, arg_1):
    """Creates `dist_fn`, `dist_fn_wrapped`, `dist_fn_args`."""
    if not isinstance(arg_1, collections.Sequence):
      raise TypeError('`model` must be `list`-like (saw: {}).'.format(
          type(arg_1).__name__))
    arg_0._dist_fn = arg_1
    arg_0._dist_fn_wrapped, arg_0._dist_fn_args = zip(*[
        _unify_call_signature(i, dist_fn)
        for i, dist_fn in enumerate(arg_1)])


def Func(arg_0, arg_1=None, arg_2='x'):
    """Creates a `tuple` of `tuple`s of dependencies.

    This function is **experimental**. That said, we encourage its use
    and ask that you report problems to `tfprobability@tensorflow.org`.

    Args:
      distribution_names: `list` of `str` or `None` names corresponding to each
        of `model` elements. (`None`s are expanding into the
        appropriate `str`.)
      leaf_name: `str` used when no maker depends on a particular
        `model` element.

    Returns:
      graph: `tuple` of `(str tuple)` pairs representing the name of each
        distribution (maker) and the names of its dependencies.

    #### Example

    ```python
    d = tfd.JointDistributionSequential([
                     tfd.Independent(tfd.Exponential(rate=[100, 120]), 1),
        lambda    e: tfd.Gamma(concentration=e[..., 0], rate=e[..., 1]),
                     tfd.Normal(loc=0, scale=2.),
        lambda n, g: tfd.Normal(loc=n, scale=g),
    ])
    d.Func()
    # ==> (
    #       ('e', ()),
    #       ('g', ('e',)),
    #       ('n', ()),
    #       ('x', ('n', 'g')),
    #     )
    ```

    """
    # This function additionally depends on:
    #   self._dist_fn_args
    #   self._dist_fn_wrapped
    # TODO(b/129008220): Robustify this procedure. Eg, handle collisions better,
    # ignore args prefixed with `_`.
    if arg_1 is None or any(arg_0._dist_fn_args):
      arg_1 = _resolve_distribution_names(
          arg_0._dist_fn_args, arg_1, arg_2)
    if len(set(arg_1)) != len(arg_1):
      raise ValueError('Distribution names must be unique: {}'.format(
          arg_1))
    if len(arg_1) != len(arg_0._dist_fn_wrapped):
      raise ValueError('Distribution names must be 1:1 with `rvs`.')
    return tuple(zip(arg_1,
                     tuple(() if arg_3 is None else arg_3 for arg_3 in arg_0._dist_fn_args)))


def Func(arg_0):
  """Decorator function for argument bounds checking.

  This decorator is meant to be used with methods that require the first
  argument to be in the support of the distribution. If `validate_args` is
  `True`, the method is wrapped with an assertion that the first argument is
  greater than or equal to `loc`, since the support of the half-Cauchy
  distribution is given by `[loc, infinity)`.

  Args:
    f: method to be decorated.

  Returns:
    Returns a decorated method that, when `validate_args` attribute of the class
    is `True`, will assert that all elements in the first argument are within
    the support of the distribution before executing the original method.
  """
  @functools.wraps(arg_0)
  def _check_arg_and_apply_f(*arg_1, **arg_2):
    arg_3 = arg_1[0]
    arg_4 = arg_1[1]
    with tf.control_dependencies([
        assert_util.assert_greater_equal(
            arg_4, arg_3.loc, message="x is not in the support of the distribution")
    ] if arg_3.validate_args else []):
      return arg_0(*arg_1, **arg_2)
  return _check_arg_and_apply_f


def Func(arg_0, arg_1, arg_2=None):
  """Visualizes sequences as TensorBoard summaries.

  Args:
    seqs: A tensor of shape [n, t, h, w, c].
    name: String name of this summary.
    num: Integer for the number of examples to visualize. Defaults to
      all examples.
  """
  arg_0 = tf.clip_by_value(arg_0, 0., 1.)
  arg_0 = tf.unstack(arg_0[:arg_2])
  arg_3 = [tf.concat(tf.unstack(seq), 1) for seq in arg_0]
  arg_3 = tf.expand_dims(tf.concat(arg_3, 0), 0)
  tf.compat.v2.summary.image(
      arg_1,
      arg_3,
      max_outputs=1,
      step=tf.compat.v1.train.get_or_create_global_step())


def Func(arg_0, arg_1, arg_2=3, arg_3="reconstruction"):
  """Visualizes the reconstruction of inputs in TensorBoard.

  Args:
    inputs: A tensor of the original inputs, of shape [batch, timesteps,
      h, w, c].
    reconstruct: A tensor of a reconstruction of inputs, of shape
      [batch, timesteps, h, w, c].
    num: Integer for the number of examples to visualize.
    name: String name of this summary.
  """
  arg_1 = tf.clip_by_value(arg_1, 0., 1.)
  arg_4 = tf.concat((arg_0[:arg_2], arg_1[:arg_2]), axis=0)
  image_summary(arg_4, arg_3)


def Func(arg_0, arg_1, arg_2="dist_params"):
  """Summarize the parameters of a distribution.

  Args:
    dist: A Distribution object with mean and standard deviation
      parameters.
    name: The name of the distribution.
    name_scope: The name scope of this summary.
  """
  with tf.compat.v1.name_scope(arg_2):
    tf.compat.v2.summary.histogram(
        arg_1="{}/{}".format(arg_1, "mean"),
        data=arg_0.mean(),
        step=tf.compat.v1.train.get_or_create_global_step())
    tf.compat.v2.summary.histogram(
        arg_1="{}/{}".format(arg_1, "stddev"),
        data=arg_0.stddev(),
        step=tf.compat.v1.train.get_or_create_global_step())


def Func(arg_0, arg_1, arg_2):
    """Runs the model to generate a distribution for a single timestep.

    This generates a batched MultivariateNormalDiag distribution using
    the output of the recurrent model at the current timestep to
    parameterize the distribution.

    Args:
      inputs: The sampled value of `z` at the previous timestep, i.e.,
        `z_{t-1}`, of shape [..., dimensions].
        `z_0` should be set to the empty matrix.
      state: A tuple containing the (hidden, cell) state.

    Returns:
      A tuple of a MultivariateNormalDiag distribution, and the state of
      the recurrent function at the end of the current timestep. The
      distribution will have event shape [dimensions], batch shape
      [...], and sample shape [sample_shape, ..., dimensions].
    """
    # In order to allow the user to pass in a single example without a batch
    # dimension, we always expand the input to at least two dimensions, then
    # fix the output shape to remove the batch dimension if necessary.
    arg_3 = arg_1.shape
    if len(arg_3) < 2:
      arg_1 = tf.reshape(arg_1, [1, -1])
    arg_4, arg_2 = arg_0.lstm_cell(arg_1, arg_2)
    arg_4 = arg_0.output_layer(arg_4)
    arg_5 = tf.concat((arg_3[:-1], tf.shape(input=arg_4)[-1:]),
                              0)
    arg_4 = tf.reshape(arg_4, arg_5)
    arg_6 = arg_4[..., :arg_0.dimensions]
    arg_7 = tf.nn.softplus(arg_4[..., arg_0.dimensions:]) + 1e-5  # keep > 0
    return tfd.MultivariateNormalDiag(arg_6=arg_6, arg_7=arg_7), arg_2


def Func(arg_0):
    """Static batch shape of models represented by this component.

    Returns:
      batch_shape: A `tf.TensorShape` giving the broadcast batch shape of
        all model parameters. This should match the batch shape of
        derived state space models, i.e.,
        `self.make_state_space_model(...).batch_shape`. It may be partially
        defined or unknown.
    """
    Func = tf.TensorShape([])
    for arg_2 in arg_0.parameters:
      Func = tf.broadcast_static_shape(
          Func, arg_2.prior.batch_shape)
    return Func


def Func(arg_0,
                             arg_1,
                             arg_2=None,
                             arg_3=None,
                             arg_4=0):
    """Instantiate this model as a Distribution over specified `num_timesteps`.

    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_state_prior: an optional `Distribution` instance overriding the
        default prior on the model's initial state. This is used in forecasting
        ("today's prior is yesterday's posterior").
      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:
      dist: a `LinearGaussianStateSpaceModel` Distribution object.
    """
    return arg_0._Func(
        arg_1=arg_1,
        param_map=arg_0._canonicalize_param_vals_as_map(arg_2),
        arg_3=arg_3,
        arg_4=arg_4)


def Func(arg_0,
                   arg_1,
                   arg_2=0,
                   arg_3=(),
                   arg_4=(),
                   arg_5=None):
    """Sample from the joint prior over model parameters and trajectories.

    Args:
      num_timesteps: Scalar `int` `Tensor` number of timesteps to model.
      initial_step: Optional scalar `int` `Tensor` specifying the starting
        timestep.
          Default value: 0.
      params_sample_shape: Number of possible worlds to sample iid from the
        parameter prior, or more generally, `Tensor` `int` shape to fill with
        iid samples.
          Default value: [] (i.e., draw a single sample and don't expand the
          shape).
      trajectories_sample_shape: For each sampled set of parameters, number
        of trajectories to sample, or more generally, `Tensor` `int` shape to
        fill with iid samples.
        Default value: [] (i.e., draw a single sample and don't expand the
          shape).
      seed: Python `int` random seed.

    Returns:
      trajectories: `float` `Tensor` of shape
        `trajectories_sample_shape + params_sample_shape + [num_timesteps, 1]`
        containing all sampled trajectories.
      param_samples: list of sampled parameter value `Tensor`s, in order
        corresponding to `self.parameters`, each of shape
        `params_sample_shape + prior.batch_shape + prior.event_shape`.
    """

    arg_5 = distributions.SeedStream(
        arg_5, salt='StructuralTimeSeries_Func')

    with tf.compat.v1.name_scope(
        'Func',
        values=[arg_1, arg_3, arg_4]):
      arg_6 = [
          p.prior.sample(arg_3, arg_5=arg_5(), name=p.name)
          for p in arg_0.parameters
      ]
      arg_7 = arg_0.make_state_space_model(
          arg_1=arg_1,
          arg_2=arg_2,
          param_vals=arg_6)
      return arg_7.sample(arg_4, arg_5=arg_5()), arg_6


def Func(arg_0, arg_1=-1, arg_2='ASCENDING', arg_3=False, arg_4=None):  # pylint: disable=unused-argument
  """Numpy implementation of `tf.argsort`."""
  if arg_2 == 'ASCENDING':
    pass
  elif arg_2 == 'DESCENDING':
    arg_0 = np.negative(arg_0)
  else:
    raise ValueError('Unrecognized direction: {}.'.format(arg_2))
  return np.argsort(arg_0, arg_1, kind='stable' if arg_3 else 'quicksort')


def Func(arg_0, arg_1=-1, arg_2='ASCENDING', arg_3=False, arg_4=None):  # pylint: disable=unused-argument
  """Numpy implementation of `tf.sort`."""
  if arg_2 == 'ASCENDING':
    pass
  elif arg_2 == 'DESCENDING':
    arg_0 = np.negative(arg_0)
  else:
    raise ValueError('Unrecognized direction: {}.'.format(arg_2))
  arg_5 = np.sort(arg_0, arg_1, kind='stable' if arg_3 else 'quicksort')
  if arg_2 == 'DESCENDING':
    return np.negative(arg_5)
  return arg_5


def Func(arg_0, arg_1="Func"):
  """Normal distribution function.

  Returns the area under the Gaussian probability density function, integrated
  from minus infinity to x:

  ```
                    1       / x
     Func(x)  = ----------  |    exp(-0.5 t**2) dt
                sqrt(2 pi)  /-inf

              = 0.5 (1 + erf(x / sqrt(2)))
              = 0.5 erfc(x / sqrt(2))
  ```

  Args:
    x: `Tensor` of type `float32`, `float64`.
    name: Python string. A name for the operation (default="Func").

  Returns:
    Func: `Tensor` with `dtype=x.dtype`.

  Raises:
    TypeError: if `x` is not floating-type.
  """

  with tf.name_scope(arg_1):
    arg_0 = tf.convert_to_tensor(value=arg_0, arg_1="x")
    if dtype_util.as_numpy_dtype(arg_0.dtype) not in [np.float32, np.float64]:
      raise TypeError(
          "x.dtype=%s is not handled, see docstring for supported types."
          % arg_0.dtype)
    return _Func(arg_0)


def Func(arg_0):
  """Implements ndtr core logic."""
  arg_1 = tf.constant(
      0.5 * np.sqrt(2.), dtype=arg_0.dtype, name="half_sqrt_2")
  arg_2 = arg_0 * arg_1
  arg_3 = tf.abs(arg_2)
  arg_4 = tf.where(
      tf.less(arg_3, arg_1), 1. + tf.math.erf(arg_2),
      tf.where(tf.greater(arg_2, 0.), 2. - tf.math.erfc(arg_3), tf.math.erfc(arg_3)))
  return 0.5 * arg_4


def Func(arg_0, arg_1="Func"):
  """The inverse of the CDF of the Normal distribution function.

  Returns x such that the area under the pdf from minus infinity to x is equal
  to p.

  A piece-wise rational approximation is done for the function.
  This is a port of the implementation in netlib.

  Args:
    p: `Tensor` of type `float32`, `float64`.
    name: Python string. A name for the operation (default="Func").

  Returns:
    x: `Tensor` with `dtype=p.dtype`.

  Raises:
    TypeError: if `p` is not floating-type.
  """

  with tf.name_scope(arg_1):
    arg_0 = tf.convert_to_tensor(value=arg_0, arg_1="p")
    if dtype_util.as_numpy_dtype(arg_0.dtype) not in [np.float32, np.float64]:
      raise TypeError(
          "p.dtype=%s is not handled, see docstring for supported types."
          % arg_0.dtype)
    return _Func(arg_0)


def Func(arg_0, arg_1=3, arg_2="Func"):
  """Log Normal distribution function.

  For details of the Normal distribution function see `ndtr`.

  This function calculates `(log o ndtr)(x)` by either calling `log(ndtr(x))` or
  using an asymptotic series. Specifically:
  - For `x > upper_segment`, use the approximation `-ndtr(-x)` based on
    `log(1-x) ~= -x, x << 1`.
  - For `lower_segment < x <= upper_segment`, use the existing `ndtr` technique
    and take a log.
  - For `x <= lower_segment`, we use the series approximation of erf to compute
    the log CDF directly.

  The `lower_segment` is set based on the precision of the input:

  ```
  lower_segment = { -20,  x.dtype=float64
                  { -10,  x.dtype=float32
  upper_segment = {   8,  x.dtype=float64
                  {   5,  x.dtype=float32
  ```

  When `x < lower_segment`, the `ndtr` asymptotic series approximation is:

  ```
     ndtr(x) = scale * (1 + sum) + R_N
     scale   = exp(-0.5 x**2) / (-x sqrt(2 pi))
     sum     = Sum{(-1)^n (2n-1)!! / (x**2)^n, n=1:N}
     R_N     = O(exp(-0.5 x**2) (2N+1)!! / |x|^{2N+3})
  ```

  where `(2n-1)!! = (2n-1) (2n-3) (2n-5) ...  (3) (1)` is a
  [double-factorial](https://en.wikipedia.org/wiki/Double_factorial).

  Args:
    x: `Tensor` of type `float32`, `float64`.
    series_order: Positive Python `integer`. Maximum depth to
      evaluate the asymptotic expansion. This is the `N` above.
    name: Python string. A name for the operation (default="Func").

  Returns:
    Func: `Tensor` with `dtype=x.dtype`.

  Raises:
    TypeError: if `x.dtype` is not handled.
    TypeError: if `series_order` is a not Python `integer.`
    ValueError:  if `series_order` is not in `[0, 30]`.
  """
  if not isinstance(arg_1, int):
    raise TypeError("series_order must be a Python integer.")
  if arg_1 < 0:
    raise ValueError("series_order must be non-negative.")
  if arg_1 > 30:
    raise ValueError("series_order must be <= 30.")

  with tf.name_scope(arg_2):
    arg_0 = tf.convert_to_tensor(value=arg_0, arg_2="x")

    if dtype_util.base_equal(arg_0.dtype, tf.float64):
      arg_3 = LOGNDTR_FLOAT64_LOWER
      arg_4 = LOGNDTR_FLOAT64_UPPER
    elif dtype_util.base_equal(arg_0.dtype, tf.float32):
      arg_3 = LOGNDTR_FLOAT32_LOWER
      arg_4 = LOGNDTR_FLOAT32_UPPER
    else:
      raise TypeError("x.dtype=%s is not supported." % arg_0.dtype)

    # The basic idea here was ported from:
    #   https://root.cern.ch/doc/v608/SpecFuncCephesInv_8cxx_source.html
    # We copy the main idea, with a few changes
    # * For x >> 1, and X ~ Normal(0, 1),
    #     Log[P[X < x]] = Log[1 - P[X < -x]] approx -P[X < -x],
    #     which extends the range of validity of this function.
    # * We use one fixed series_order for all of 'x', rather than adaptive.
    # * Our docstring properly reflects that this is an asymptotic series, not a
    #   Taylor series. We also provided a correct bound on the remainder.
    # * We need to use the max/min in the _Func_lower arg to avoid nan when
    #   x=0. This happens even though the branch is unchosen because when x=0
    #   the gradient of a select involves the calculation 1*dy+0*(-inf)=nan
    #   regardless of whether dy is finite. Note that the minimum is a NOP if
    #   the branch is chosen.
    return tf.where(
        tf.greater(arg_0, arg_4),
        -_ndtr(-arg_0),  # log(1-x) ~= -x, x << 1
        tf.where(
            tf.greater(arg_0, arg_3),
            tf.math.log(_ndtr(tf.maximum(arg_0, arg_3))),
            _Func_lower(tf.minimum(arg_0, arg_3), arg_1)))


def Func(arg_0, arg_1):
  """Calculates the asymptotic series used in log_ndtr."""
  arg_2 = dtype_util.as_numpy_dtype(arg_0.dtype)
  if arg_1 <= 0:
    return arg_2(1)
  arg_3 = tf.square(arg_0)
  arg_4 = tf.zeros_like(arg_0)
  arg_5 = tf.zeros_like(arg_0)
  arg_6 = arg_3  # Start with x^{2*1} = x^{2*n} with n = 1.
  for arg_7 in range(1, arg_1 + 1):
    arg_8 = arg_2(_double_factorial(2 * arg_7 - 1)) / arg_6
    if arg_7 % 2:
      arg_5 += arg_8
    else:
      arg_4 += arg_8
    arg_6 *= arg_3
  return 1. + arg_4 - arg_5


def Func(arg_0, arg_1="Func"):
  """Log Laplace distribution function.

  This function calculates `Log[L(x)]`, where `L(x)` is the cumulative
  distribution function of the Laplace distribution, i.e.

  ```L(x) := 0.5 * int_{-infty}^x e^{-|t|} dt```

  For numerical accuracy, `L(x)` is computed in different ways depending on `x`,

  ```
  x <= 0:
    Log[L(x)] = Log[0.5] + x, which is exact

  0 < x:
    Log[L(x)] = Log[1 - 0.5 * e^{-x}], which is exact
  ```

  Args:
    x: `Tensor` of type `float32`, `float64`.
    name: Python string. A name for the operation (default="log_ndtr").

  Returns:
    `Tensor` with `dtype=x.dtype`.

  Raises:
    TypeError: if `x.dtype` is not handled.
  """

  with tf.name_scope(arg_1):
    arg_0 = tf.convert_to_tensor(value=arg_0, arg_1="x")

    # For x < 0, L(x) = 0.5 * exp{x} exactly, so Log[L(x)] = log(0.5) + x.
    arg_2 = -np.log(2.) + arg_0

    # safe_exp_neg_x = exp{-x} for x > 0, but is
    # bounded above by 1, which avoids
    #   log[1 - 1] = -inf for x = log(1/2), AND
    #   exp{-x} --> inf, for x << -1
    arg_3 = tf.exp(-tf.abs(arg_0))

    # log1p(z) = log(1 + z) approx z for |z| << 1. This approxmation is used
    # internally by log1p, rather than being done explicitly here.
    arg_4 = tf.math.log1p(-0.5 * arg_3)

    return tf.where(arg_0 < 0., arg_2, arg_4)


def Func(arg_0, arg_1, arg_2, arg_3):
  """Joint log probability function."""
  arg_4 = (1. / tf.reduce_mean(input_tensor=arg_0))
  arg_5 = tfd.Exponential(rate=arg_4)

  arg_6 = tfd.Uniform()

  arg_7 = tf.gather(
      [arg_1, arg_2],
      indices=tf.cast(
          arg_3 * tf.cast(tf.size(input=arg_0), dtype=tf.float32) <= tf.cast(
              tf.range(tf.size(input=arg_0)), dtype=tf.float32),
          dtype=tf.int32))
  arg_8 = tfd.Poisson(rate=arg_7)

  return (arg_5.log_prob(arg_1) + arg_5.log_prob(arg_2) +
          arg_6.log_prob(arg_3) +
          tf.reduce_sum(input_tensor=arg_8.log_prob(arg_0)))


def Func(
    arg_0=arg_1(3e3),
    arg_2=arg_1(3e3),
    arg_3=3):
  """Runs HMC on the text-messages unnormalized posterior."""

  if not tf.executing_eagerly():
    tf.compat.v1.reset_default_graph()

  # Build a static, pretend dataset.
  arg_4 = tf.cast(
      tf.concat(
          [tfd.Poisson(rate=15.).sample(43),
           tfd.Poisson(rate=25.).sample(31)],
          axis=0),
      dtype=tf.float32)
  if tf.executing_eagerly():
    arg_4 = arg_4.numpy()
  else:
    with tf.compat.v1.Session():
      arg_4 = arg_4.eval()

  # Define a closure over our joint_log_prob.
  def unnormalized_log_posterior(arg_5, arg_6, arg_7):
    return text_messages_joint_log_prob(arg_4, arg_5, arg_6, arg_7)

  if tf.executing_eagerly():
    arg_8 = tf.function(tfp.mcmc.sample_chain)
  else:
    arg_8 = tfp.mcmc.sample_chain

  # Initialize the step_size. (It will be automatically adapted.)
  arg_9 = tf.compat.v2.Variable(
      name='step_size',
      initial_value=tf.constant(0.05, dtype=tf.float32),
      trainable=False)

  def computation():
    """The benchmark computation."""

    arg_10 = [
        tf.constant(arg_4.mean(), name='init_lambda1'),
        tf.constant(arg_4.mean(), name='init_lambda2'),
        tf.constant(0.5, name='init_tau'),
    ]

    arg_11 = [
        tfp.bijectors.Exp(),       # Maps a positive real to R.
        tfp.bijectors.Exp(),       # Maps a positive real to R.
        tfp.bijectors.Sigmoid(),   # Maps [0,1] to R.
    ]

    arg_12, arg_13 = arg_8(
        arg_0=arg_0,
        arg_2=arg_2,
        current_state=arg_10,
        kernel=tfp.mcmc.TransformedTransitionKernel(
            inner_kernel=tfp.mcmc.HamiltonianMonteCarlo(
                target_log_prob_fn=unnormalized_log_posterior,
                arg_3=arg_3,
                arg_9=arg_9,
                step_size_update_fn=
                tfp.mcmc.make_simple_step_size_update_policy(arg_2),
                state_gradients_are_stopped=True),
            bijector=arg_11))

    return arg_13.inner_results.is_accepted

  # Let's force evaluation of graph to ensure build time is not part of our time
  # trial.
  arg_14 = computation()
  if not tf.executing_eagerly():
    arg_15 = tf.compat.v1.Session()
    arg_15.run(tf.compat.v1.global_variables_initializer())
    arg_15.run(arg_14)

  arg_16 = time.time()
  if tf.executing_eagerly():
    arg_17 = computation()
  else:
    arg_17 = arg_15.run(arg_14)
  arg_18 = time.time() - arg_16

  arg_19 = np.sum(arg_17)
  arg_20 = np.float32(arg_19) / np.float32(arg_0)

  return dict(
      iters=(arg_0 + arg_2) * arg_3,
      extras={'acceptance_rate': arg_20},
      arg_18=arg_18)


def Func(arg_0, arg_1=None):
    """Compute the marginal of this GP over function values at `index_points`.

    Args:
      index_points: `float` `Tensor` representing finite (batch of) vector(s) of
        points in the index set over which the GP is defined. Shape has the form
        `[b1, ..., bB, e, f1, ..., fF]` where `F` is the number of feature
        dimensions and must equal `kernel.feature_ndims` and `e` is the number
        (size) of index points in each batch. Ultimately this distribution
        corresponds to a `e`-dimensional multivariate normal. The batch shape
        must be broadcastable with `kernel.batch_shape` and any batch dims
        yielded by `mean_fn`.

    Returns:
      marginal: a `Normal` or `MultivariateNormalLinearOperator` distribution,
        according to whether `index_points` consists of one or many index
        points, respectively.
    """
    with arg_0._name_scope('Func'):
      # TODO(cgs): consider caching the result here, keyed on `index_points`.
      arg_1 = arg_0._get_index_points(arg_1)
      arg_2 = arg_0._compute_covariance(arg_1)
      arg_3 = arg_0._mean_fn(arg_1)
      # If we're sure the number of index points is 1, we can just construct a
      # scalar Normal. This has computational benefits and supports things like
      # CDF that aren't otherwise straightforward to provide.
      if arg_0._is_univariate_marginal(arg_1):
        arg_4 = tf.sqrt(arg_2)
        # `loc` has a trailing 1 in the shape; squeeze it.
        arg_3 = tf.squeeze(arg_3, axis=-1)
        return normal.Normal(
            arg_3=arg_3,
            arg_4=arg_4,
            validate_args=arg_0._validate_args,
            allow_nan_stats=arg_0._allow_nan_stats,
            name='marginal_distribution')
      else:
        arg_4 = tf.linalg.LinearOperatorLowerTriangular(
            tf.linalg.cholesky(_add_diagonal_shift(arg_2, arg_0.jitter)),
            is_non_singular=True,
            name='GaussianProcessScaleLinearOperator')
        return mvn_linear_operator.MultivariateNormalLinearOperator(
            arg_3=arg_3,
            arg_4=arg_4,
            validate_args=arg_0._validate_args,
            allow_nan_stats=arg_0._allow_nan_stats,
            name='marginal_distribution')


def Func(arg_0, arg_1=None):
    """Return `index_points` if not None, else `self._index_points`.

    Args:
      index_points: if given, this is what is returned; else,
      `self._index_points`

    Returns:
      index_points: the given arg, if not None, else the class member
      `self._index_points`.

    Rases:
      ValueError: if `index_points` and `self._index_points` are both `None`.
    """
    if arg_0._index_points is None and arg_1 is None:
      raise ValueError(
          'This GaussianProcess instance was not instantiated with a value for '
          'index_points. One must therefore be provided when calling sample, '
          'log_prob, and other such methods. In particular, one can\'t compute '
          'KL divergences to/from an instance of `GaussianProccess` with '
          'unspecified `index_points` directly. Instead, use the '
          '`get_marginal_distribution` function, which takes `index_points` as '
          'an argument and returns a `Normal` or '
          '`MultivariateNormalLinearOperator` instance, whose KL can be '
          'computed.')
    return arg_1 if arg_1 is not None else arg_0._index_points


def Func(arg_0,
                   arg_1=2,
                   arg_2=None,
                   arg_3=arg_4.float32):
  """Creates an stacked IAF bijector.

  This bijector operates on vector-valued events.

  Args:
    total_event_size: Number of dimensions to operate over.
    num_hidden_layers: How many hidden layers to use in each IAF.
    seed: Random seed for the initializers.
    dtype: DType for the variables.

  Returns:
    bijector: The created bijector.
  """

  arg_2 = tfd.SeedStream(arg_2, 'Func')

  def make_iaf():
    """Create an IAF."""
    arg_6 = arg_4.compat.v2.keras.initializers.VarianceScaling(
        2 * 0.01, arg_2=arg_2() % (2**31 - 1))

    arg_7 = tfb.AutoregressiveLayer(
        params=2,
        event_shape=[arg_0],
        hidden_units=[arg_0] * arg_1,
        activation=arg_4.nn.elu,
        kernel_initializer=arg_6,
        arg_3=arg_3)

    def shift_and_scale(arg_8):
      # TODO(siege): Something is losing the static shape.
      arg_8.set_shape(
          arg_8.shape.merge_with([None] * (arg_8.shape.ndims - 1) + [arg_0]))
      return arg_4.unstack(arg_7(arg_8), num=2, axis=-1)

    return tfb.Invert(tfb.MaskedAutoregressiveFlow(shift_and_scale))

  def make_swap():
    """Create an swap."""
    arg_9 = list(reversed(range(arg_0)))
    return tfb.Permute(arg_9)

  arg_10 = make_iaf()
  arg_10 = make_swap()(arg_10)
  arg_10 = make_iaf()(arg_10)
  arg_10 = make_swap()(arg_10)
  arg_10 = make_iaf()(arg_10)
  arg_10 = make_swap()(arg_10)

  return arg_10


def Func(arg_0, arg_1, arg_2):
    """Runs one iteration of NeuTra.

    Args:
      current_state: `Tensor` or Python `list` of `Tensor`s representing the
        current state(s) of the Markov chain(s). The first `r` dimensions index
        independent chains, `r = tf.rank(target_log_prob_fn(*current_state))`.
      previous_kernel_results: `collections.namedtuple` containing `Tensor`s
        representing values from previous calls to this function (or from the
        `bootstrap_results` function.)

    Returns:
      next_state: Tensor or Python list of `Tensor`s representing the state(s)
        of the Markov chain(s) after taking exactly one step. Has same type and
        shape as `current_state`.
      kernel_results: `collections.namedtuple` of internal calculations used to
        advance the chain.
    """

    @tfp.mcmc.internal.util.make_innermost_setter
    def set_num_leapfrog_steps(arg_3, arg_4):
      return arg_3._replace(
          accepted_results=arg_3.accepted_results._replace(
              arg_4=arg_4))

    arg_5 = arg_2.new_step_size
    arg_2 = set_num_leapfrog_steps(
        arg_2, arg_0._num_leapfrog_steps(arg_5))

    arg_6, arg_3 = arg_0._kernel.Func(
        arg_0._flatten_state(arg_1), arg_2)
    return arg_0._unflatten_state(arg_6), arg_3


def Func(arg_0, arg_1):
    """Trains the bijector and creates initial `previous_kernel_results`.

    The supplied `state` is only used to determine the number of chains to run
    in parallel_iterations

    Args:
      state: `Tensor` or Python `list` of `Tensor`s representing the initial
        state(s) of the Markov chain(s). The first `r` dimensions index
        independent chains, `r = tf.rank(target_log_prob_fn(*state))`.

    Returns:
      kernel_results: Instance of
        `UncalibratedHamiltonianMonteCarloKernelResults` inside
        `MetropolisHastingsResults` inside `TransformedTransitionKernelResults`
        inside `SimpleStepSizeAdaptationResults`.
    """

    def loss():
      arg_2 = arg_0._flattened_variational_distribution()
      # TODO(siege): How to seed this?
      arg_3 = arg_2.sample(arg_0.train_batch_size)
      return tf.reduce_mean(
          input_tensor=arg_2.log_prob(arg_3) -
          arg_0._flattened_target_log_prob(arg_3),
          axis=-1)

    arg_4 = tf.convert_to_tensor(value=arg_0.learning_rate, arg_5=arg_0._dtype)
    arg_5 = arg_4.dtype

    arg_6 = tf.compat.v2.optimizers.schedules.PiecewiseConstantDecay(
        list(arg_0.num_train_steps *
             np.array([0.2, 0.8]).astype(arg_5.as_numpy_dtype())),
        [arg_4, arg_4 * 0.1, arg_4 * 0.01])

    arg_7 = tf.compat.v2.optimizers.Adam(arg_6)

    @tf.function(autograph=False)
    def train_step():
      with tf.GradientTape() as tape:
        arg_8 = loss()
      arg_9 = tape.watched_variables()
      arg_10 = tape.gradient(arg_8, arg_9)
      arg_11 = list(zip(arg_10, arg_9))
      arg_7.apply_gradients(arg_11)
      return arg_8

    for arg_12 in range(arg_0.num_train_steps):
      arg_8 = train_step()
      tf.debugging.assert_all_finite(
          arg_8, 'NeuTra loss is NaN at step {}'.format(arg_12))
      if arg_0.train_debug_fn:
        # pylint: disable=not-callable
        arg_0.train_debug_fn(arg_0, arg_12, arg_8)

    arg_13 = tf.nest.flatten(arg_1)
    arg_14 = tf.nest.flatten(arg_0.state_shape)
    arg_15 = tf.shape(input=arg_13[0])[:-arg_14[0].ndims]

    return arg_0._kernel.Func(
        arg_0._flattened_variational_distribution().sample(
            arg_15, seed=arg_0.seed))


def Func(arg_0, arg_1):
  """Convenience function analogous to tf.squared_difference."""
  arg_2 = arg_0 - arg_1
  return arg_2[..., tf.newaxis, :] * arg_2[..., tf.newaxis]


def Func(arg_0, arg_1):
    """Performs distributional transform of the mixture samples.

    Distributional transform removes the parameters from samples of a
    multivariate distribution by applying conditional CDFs:
      (F(x_1), F(x_2 | x1_), ..., F(x_d | x_1, ..., x_d-1))
    (the indexing is over the "flattened" event dimensions).
    The result is a sample of product of Uniform[0, 1] distributions.

    We assume that the components are factorized, so the conditional CDFs become
      F(x_i | x_1, ..., x_i-1) = sum_k w_i^k F_k (x_i),
    where w_i^k is the posterior mixture weight: for i > 0
      w_i^k = w_k prob_k(x_1, ..., x_i-1) / sum_k' w_k' prob_k'(x_1, ..., x_i-1)
    and w_0^k = w_k is the mixture probability of the k-th component.

    Arguments:
      x: Sample of mixture distribution

    Returns:
      Result of the distributional transform
    """

    if tensorshape_util.rank(arg_1.shape) is None:
      # tf.nn.softmax raises an error when applied to inputs of undefined rank.
      raise ValueError("Distributional transform does not support inputs of "
                       "undefined rank.")

    # Obtain factorized components distribution and assert that it's
    # a scalar distribution.
    if isinstance(arg_0._components_distribution, independent.Independent):
      arg_2 = arg_0._components_distribution.distribution
    else:
      arg_2 = arg_0._components_distribution

    with tf.control_dependencies([
        assert_util.assert_equal(
            arg_2.is_scalar_event(),
            True,
            message="`univariate_components` must have scalar event")
    ]):
      arg_3 = arg_0._pad_sample_dims(arg_1)  # [S, B, 1, E]
      arg_4 = arg_2.log_prob(arg_3)  # [S, B, k, E]
      arg_5 = arg_2.cdf(arg_3)  # [S, B, k, E]

      # log prob_k (x_1, ..., x_i-1)
      arg_6 = tf.reshape(
          tf.math.cumsum(
              # [S*prod(B)*k, prod(E)]
              tf.reshape(arg_4, [-1, arg_0._event_size]),
              exclusive=True,
              axis=-1),
          tf.shape(input=arg_4))  # [S, B, k, E]

      arg_7 = distribution_utils.pad_mixture_dimensions(
          arg_0.mixture_distribution.logits, arg_0, arg_0.mixture_distribution,
          arg_0._event_ndims)  # [B, k, 1]

      # Logits of the posterior weights: log w_k + log prob_k (x_1, ..., x_i-1)
      arg_8 = arg_7 + arg_6

      arg_9 = tensorshape_util.rank(arg_1.shape) - arg_0._event_ndims
      arg_10 = tf.nn.softmax(arg_8,
                                          axis=arg_9)
      return tf.reduce_sum(
          input_tensor=arg_10 * arg_5, axis=arg_9)


def Func(
    arg_0, arg_1, arg_2, arg_3):
  """Utility method to decompose a joint posterior into components.

  Args:
    model: `tfp.sts.Sum` instance defining an additive STS model.
    posterior_means: float `Tensor` of shape `concat(
      [[num_posterior_draws], batch_shape, num_timesteps, latent_size])`
      representing the posterior mean over latents in an
      `AdditiveStateSpaceModel`.
    posterior_covs: float `Tensor` of shape `concat(
      [[num_posterior_draws], batch_shape, num_timesteps,
      latent_size, latent_size])`
      representing the posterior marginal covariances over latents in an
      `AdditiveStateSpaceModel`.
    parameter_samples: Python `list` of `Tensors` representing posterior
      samples of model parameters, with shapes `[concat([
      [num_posterior_draws], param.prior.batch_shape,
      param.prior.event_shape]) for param in model.parameters]`. This may
      optionally also be a map (Python `dict`) of parameter names to
      `Tensor` values.

  Returns:
    component_dists: A `collections.OrderedDict` instance mapping
      component StructuralTimeSeries instances (elements of `model.components`)
      to `tfd.Distribution` instances representing the posterior marginal
      distributions on the process modeled by each component. Each distribution
      has batch shape matching that of `posterior_means`/`posterior_covs`, and
      event shape of `[num_timesteps]`.
  """

  try:
    arg_0.components
  except AttributeError:
    raise ValueError('Model decomposed into components must be an instance of'
                     '`tfp.sts.Sum` (passed model {})'.format(arg_0))

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

    # Extract the component means/covs from the joint latent posterior.
    arg_4 = [arg_10.latent_size for arg_10 in arg_0.components]
    arg_5 = tf.split(arg_1, arg_4, axis=-1)
    arg_6 = _split_covariance_into_marginals(
        arg_2, arg_4)

    # Instantiate per-component state space models, and use them to push the
    # posterior means/covs through the observation model for each component.
    arg_7 = dist_util.prefer_static_value(
        tf.shape(input=arg_1))[-2]
    arg_8 = arg_0.make_component_state_space_models(
        arg_7=arg_7,
        param_vals=arg_3)
    arg_9 = collections.OrderedDict()
    for (arg_10, arg_11,
         arg_12, arg_13) in zip(arg_0.components, arg_8,
                                               arg_5, arg_6):
      arg_14, arg_15 = (
          arg_11.latents_to_observations(
              latent_means=arg_12,
              latent_covs=arg_13))

      # Using the observation means and covs, build a mixture distribution
      # that integrates over the posterior draws.
      arg_9[arg_10] = sts_util.mix_over_posterior_draws(
          means=arg_14[..., 0],
          variances=arg_15[..., 0, 0])
  return arg_9


def Func(arg_0, arg_1, arg_2):
  """Decompose a forecast distribution into contributions from each component.

  Args:
    model: An instance of `tfp.sts.Sum` representing a structural time series
      model.
    forecast_dist: A `Distribution` instance returned by `tfp.sts.forecast()`.
      (specifically, must be a `tfd.MixtureSameFamily` over a
      `tfd.LinearGaussianStateSpaceModel` parameterized by posterior samples).
    parameter_samples: Python `list` of `Tensors` representing posterior samples
      of model parameters, with shapes `[concat([[num_posterior_draws],
      param.prior.batch_shape, param.prior.event_shape]) for param in
      model.parameters]`. This may optionally also be a map (Python `dict`) of
      parameter names to `Tensor` values.
  Returns:
    component_forecasts: A `collections.OrderedDict` instance mapping
      component StructuralTimeSeries instances (elements of `model.components`)
      to `tfd.Distribution` instances representing the marginal forecast for
      each component. Each distribution has batch and event shape matching
      `forecast_dist` (specifically, the event shape is
      `[num_steps_forecast]`).

  #### Examples

  Suppose we've built a model, fit it to data, and constructed a forecast
  distribution:

  ```python
    day_of_week = tfp.sts.Seasonal(
        num_seasons=7,
        observed_time_series=observed_time_series,
        name='day_of_week')
    local_linear_trend = tfp.sts.LocalLinearTrend(
        observed_time_series=observed_time_series,
        name='local_linear_trend')
    model = tfp.sts.Sum(components=[day_of_week, local_linear_trend],
                        observed_time_series=observed_time_series)

    num_steps_forecast = 50
    samples, kernel_results = tfp.sts.fit_with_hmc(model, observed_time_series)
    forecast_dist = tfp.sts.forecast(model, observed_time_series,
                                 parameter_samples=samples,
                                 num_steps_forecast=num_steps_forecast)
  ```

  To extract the forecast for individual components, pass the forecast
  distribution into `Funcs`:

  ```python
    component_forecasts = Func(
      model, forecast_dist, samples)

    # Component mean and stddev have shape `[num_steps_forecast]`.
    day_of_week_effect_mean = forecast_components[day_of_week].mean()
    day_of_week_effect_stddev = forecast_components[day_of_week].stddev()
  ```

  Using the component forecasts, we can visualize the uncertainty for each
  component:

  ```
  from matplotlib import pylab as plt
  num_components = len(component_forecasts)
  xs = np.arange(num_steps_forecast)
  fig = plt.figure(figsize=(12, 3 * num_components))
  for i, (component, component_dist) in enumerate(component_forecasts.items()):

    # If in graph mode, replace `.numpy()` with `.eval()` or `sess.run()`.
    component_mean = component_dist.mean().numpy()
    component_stddev = component_dist.stddev().numpy()

    ax = fig.add_subplot(num_components, 1, 1 + i)
    ax.plot(xs, component_mean, lw=2)
    ax.fill_between(xs,
                    component_mean - 2 * component_stddev,
                    component_mean + 2 * component_stddev,
                    alpha=0.5)
    ax.set_title(component.name)
  ```

  """

  with tf.compat.v1.name_scope('Func'):
    try:
      arg_3 = arg_1.components_distribution
      arg_4, arg_5 = arg_3._joint_mean()  # pylint: disable=protected-access
      arg_6, arg_5 = arg_3._joint_covariances()  # pylint: disable=protected-access
    except AttributeError as e:
      raise ValueError(
          'Forecast distribution must be a MixtureSameFamily of'
          'LinearGaussianStateSpaceModel distributions, such as returned by'
          '`tfp.sts.forecast()`. (saw exception: {})'.format(e))

    # Since `parameter_samples` will have sample shape `[num_posterior_draws]`,
    # we need to move the `num_posterior_draws` dimension of the forecast
    # moments from the trailing batch dimension, where it's currently put by
    # `sts.forecast`, back to the leading (sample shape) dimension.
    arg_4 = dist_util.move_dimension(
        arg_4, source_idx=-3, dest_idx=0)
    arg_6 = dist_util.move_dimension(
        arg_6, source_idx=-4, dest_idx=0)

    return _decompose_from_posterior_marginals(
        arg_0, arg_4, arg_6, arg_2)


def Func(arg_0):
    """Get tensor that the random variable corresponds to."""
    if arg_0._Func is None:
      try:
        arg_0._Func = arg_0.distribution.sample(arg_0.sample_shape_tensor())
      except NotImplementedError:
        raise NotImplementedError(
            "sample is not implemented for {0}. You must either pass in the "
            "Func argument or implement sample for {0}."
            .format(arg_0.distribution.__class__.__name__))
    return arg_0._Func


def Func(arg_0, arg_1=None, arg_2=None):
    """In a session, computes and returns the value of this random variable.

    This is not a graph construction method, it does not add ops to the graph.

    This convenience method requires a session where the graph
    containing this variable has been launched. If no session is
    passed, the default session is used.

    Args:
      session: tf.BaseSession.
        The `tf.Session` to use to Funcuate this random variable. If
        none, the default session is used.
      feed_dict: dict.
        A dictionary that maps `tf.Tensor` objects to feed values. See
        `tf.Session.run()` for a description of the valid feed values.

    Returns:
      Value of the random variable.

    #### Examples

    ```python
    x = Normal(0.0, 1.0)
    with tf.Session() as sess:
      # Usage passing the session explicitly.
      print(x.Func(sess))
      # Usage with the default session.  The 'with' block
      # above makes 'sess' the default session.
      print(x.Func())
    ```
    """
    return arg_0.value.Func(arg_1=arg_1, arg_2=arg_2)


def Func(arg_0):
    """Value as NumPy array, only available for TF Eager."""
    if not isinstance(arg_0.value, ops.EagerTensor):
      raise NotImplementedError("value argument must be a EagerTensor.")

    return arg_0.value.Func()


def Func(arg_0, arg_1, arg_2, arg_3):
  """Posterior Normal distribution with conjugate prior on the mean.

  This model assumes that `n` observations (with sum `s`) come from a
  Normal with unknown mean `loc` (described by the Normal `prior`)
  and known variance `scale**2`. The "known scale posterior" is
  the distribution of the unknown `loc`.

  Accepts a prior Normal distribution object, having parameters
  `loc0` and `scale0`, as well as known `scale` values of the predictive
  distribution(s) (also assumed Normal),
  and statistical estimates `s` (the sum(s) of the observations) and
  `n` (the number(s) of observations).

  Returns a posterior (also Normal) distribution object, with parameters
  `(loc', scale'**2)`, where:

  ```
  mu ~ N(mu', sigma'**2)
  sigma'**2 = 1/(1/sigma0**2 + n/sigma**2),
  mu' = (mu0/sigma0**2 + s/sigma**2) * sigma'**2.
  ```

  Distribution parameters from `prior`, as well as `scale`, `s`, and `n`.
  will broadcast in the case of multidimensional sets of parameters.

  Args:
    prior: `Normal` object of type `dtype`:
      the prior distribution having parameters `(loc0, scale0)`.
    scale: tensor of type `dtype`, taking values `scale > 0`.
      The known stddev parameter(s).
    s: Tensor of type `dtype`. The sum(s) of observations.
    n: Tensor of type `int`. The number(s) of observations.

  Returns:
    A new Normal posterior distribution object for the unknown observation
    mean `loc`.

  Raises:
    TypeError: if dtype of `s` does not match `dtype`, or `prior` is not a
      Normal object.
  """
  if not isinstance(arg_0, normal.Normal):
    raise TypeError("Expected prior to be an instance of type Normal")

  if arg_2.dtype != arg_0.dtype:
    raise TypeError(
        "Observation sum s.dtype does not match prior dtype: %s vs. %s"
        % (arg_2.dtype, arg_0.dtype))

  arg_3 = tf.cast(arg_3, arg_0.dtype)
  arg_4 = tf.square(arg_0.scale)
  arg_5 = tf.square(arg_1)
  arg_6 = 1.0/(1/arg_4 + arg_3/arg_5)
  return normal.Normal(
      loc=(arg_0.loc / arg_4 + arg_2 / arg_5) * arg_6,
      arg_1=tf.sqrt(arg_6))


def Func(arg_0,
                              arg_1=False,
                              arg_2=arg_3.nn.relu,
                              arg_6=None,
                              *arg_7,  # pylint: disable=keyword-arg-before-vararg
                              **arg_8):
  """Build a scale-and-shift function using a multi-layer neural network.

  This will be wrapped in a make_template to ensure the variables are only
  created once. It takes the `d`-dimensional input x[0:d] and returns the `D-d`
  dimensional outputs `loc` ("mu") and `log_scale` ("alpha").

  The default template does not support conditioning and will raise an
  exception if `condition_kwargs` are passed to it. To use conditioning in
  real nvp bijector, implement a conditioned shift/scale template that
  handles the `condition_kwargs`.

  Arguments:
    hidden_layers: Python `list`-like of non-negative integer, scalars
      indicating the number of units in each hidden layer. Default: `[512, 512].
    shift_only: Python `bool` indicating if only the `shift` term shall be
      computed (i.e. NICE bijector). Default: `False`.
    activation: Activation function (callable). Explicitly setting to `None`
      implies a linear activation.
    name: A name for ops managed by this function. Default:
      "Func".
    *args: `tf.layers.dense` arguments.
    **kwargs: `tf.layers.dense` keyword arguments.

  Returns:
    shift: `Float`-like `Tensor` of shift terms ("mu" in
      [Papamakarios et al.  (2016)][1]).
    log_scale: `Float`-like `Tensor` of log(scale) terms ("alpha" in
      [Papamakarios et al. (2016)][1]).

  Raises:
    NotImplementedError: if rightmost dimension of `inputs` is unknown prior to
      graph execution, or if `condition_kwargs` is not empty.

  #### References

  [1]: George Papamakarios, Theo Pavlakou, and Iain Murray. Masked
       Autoregressive Flow for Density Estimation. In _Neural Information
       Processing Systems_, 2017. https://arxiv.org/abs/1705.07057
  """

  with arg_3.compat.v2.name_scope(arg_6 or "Func"):

    def _fn(arg_9, arg_10, **arg_11):
      """Fully connected MLP parameterized via `real_nvp_template`."""
      if arg_11:
        raise NotImplementedError(
            "Conditioning not implemented in the default template.")

      if tensorshape_util.rank(arg_9.shape) == 1:
        arg_9 = arg_9[arg_3.newaxis, ...]
        arg_12 = lambda arg_9: arg_9[0]
      else:
        arg_12 = lambda arg_9: arg_9
      for arg_13 in arg_0:
        arg_9 = arg_3.compat.v1.layers.dense(
            inputs=arg_9,
            arg_13=arg_13,
            arg_2=arg_2,
            *arg_7,  # pylint: disable=keyword-arg-before-vararg
            **arg_8)
      arg_9 = arg_3.compat.v1.layers.dense(
          inputs=arg_9,
          arg_13=(1 if arg_1 else 2) * arg_10,
          arg_2=None,
          *arg_7,  # pylint: disable=keyword-arg-before-vararg
          **arg_8)
      if arg_1:
        return arg_12(arg_9), None
      arg_14, arg_15 = arg_3.split(arg_9, 2, axis=-1)
      return arg_12(arg_14), arg_12(arg_15)

    return arg_3.compat.v1.make_template("Func", _fn)


def Func(arg_0, arg_1, arg_2, arg_3):
  """Returns a batch of points chosen uniformly from the unit hypersphere."""
  # This works because the Gaussian distribution is spherically symmetric.
  # raw shape: shape + [dimension]
  arg_4 = normal.Normal(
      loc=dtype_util.as_numpy_dtype(arg_2)(0),
      scale=dtype_util.as_numpy_dtype(arg_2)(1)).sample(
          tf.concat([arg_1, [arg_0]], axis=0), arg_3=arg_3())
  arg_5 = arg_4 / tf.norm(tensor=arg_4, ord=2, axis=-1)[..., tf.newaxis]
  return arg_5


def Func(arg_0, arg_1='log_normalization'):
    """Returns the log normalization of an LKJ distribution.

    Args:
      name: Python `str` name prefixed to Ops created by this function.

    Returns:
      log_z: A Tensor of the same shape and dtype as `concentration`, containing
        the corresponding log normalizers.
    """
    # The formula is from D. Lewandowski et al [1], p. 1999, from the
    # proof that eqs 16 and 17 are equivalent.
    with tf.name_scope(arg_1 or 'log_normalization_lkj'):
      arg_2 = np.log(np.pi)
      arg_3 = tf.zeros_like(arg_0.concentration)
      for arg_4 in range(1, arg_0.dimension):
        arg_3 += arg_2 * (arg_4 / 2.)
        arg_3 += tf.math.lgamma(arg_0.concentration +
                              (arg_0.dimension - 1 - arg_4) / 2.)
        arg_3 -= tf.math.lgamma(arg_0.concentration + (arg_0.dimension - 1) / 2.)
      return arg_3


def Func(arg_0, arg_1=None):
  """Returns explict dtype from `args_list` if exists, else preferred_dtype."""
  arg_2 = None
  arg_1 = (None if arg_1 is None
                     else tf.as_dtype(arg_1))
  for arg_3 in tf.nest.flatten(arg_0):
    if hasattr(arg_3, 'dtype'):
      arg_4 = tf.as_dtype(arg_3.dtype)
    else:
      continue
    if arg_2 is None:
      arg_2 = arg_4
    elif arg_2 != arg_4:
      raise TypeError('Found incompatible dtypes, {} and {}.'.format(arg_2, arg_4))
  if arg_2 is None and arg_1 is None:
    return None
  return (arg_1 if arg_2 is None else arg_2).as_numpy_dtype


def Func(arg_0):
  """Factory for implementing summary statistics, eg, mean, stddev, mode."""
  def _fn(arg_1, **arg_2):
    """Implements summary statistic, eg, mean, stddev, mode."""
    arg_3 = getattr(arg_1.distribution, arg_0)(**arg_2)
    arg_4 = prefer_static.concat([
        arg_1.distribution.batch_shape_tensor(),
        prefer_static.ones(prefer_static.rank_from_shape(arg_1.sample_shape),
                           dtype=arg_1.sample_shape.dtype),
        arg_1.distribution.event_shape_tensor(),
    ], axis=0)
    arg_3 = tf.reshape(arg_3, arg_4=arg_4)
    arg_4 = prefer_static.concat([
        arg_1.distribution.batch_shape_tensor(),
        arg_1.sample_shape,
        arg_1.distribution.event_shape_tensor(),
    ], axis=0)
    return tf.broadcast_to(arg_3, arg_4)
  return _fn


def Func(arg_0,
                          arg_1=0.,
                          arg_2=None,
                          arg_3=None):
  """Estimate a lower bound on effective sample size for each independent chain.

  Roughly speaking, "effective sample size" (ESS) is the size of an iid sample
  with the same variance as `state`.

  More precisely, given a stationary sequence of possibly correlated random
  variables `X_1, X_2,...,X_N`, each identically distributed ESS is the number
  such that

  ```Variance{ N**-1 * Sum{X_i} } = ESS**-1 * Variance{ X_1 }.```

  If the sequence is uncorrelated, `ESS = N`.  In general, one should expect
  `ESS <= N`, with more highly correlated sequences having smaller `ESS`.

  Args:
    states:  `Tensor` or list of `Tensor` objects.  Dimension zero should index
      identically distributed states.
    filter_threshold:  `Tensor` or list of `Tensor` objects.
      Must broadcast with `state`.  The auto-correlation sequence is truncated
      after the first appearance of a term less than `filter_threshold`.
      Setting to `None` means we use no threshold filter.  Since `|R_k| <= 1`,
      setting to any number less than `-1` has the same effect.
    filter_beyond_lag:  `Tensor` or list of `Tensor` objects.  Must be
      `int`-like and scalar valued.  The auto-correlation sequence is truncated
      to this length.  Setting to `None` means we do not filter based on number
      of lags.
    name:  `String` name to prepend to created ops.

  Returns:
    ess:  `Tensor` or list of `Tensor` objects.  The effective sample size of
      each component of `states`.  Shape will be `states.shape[1:]`.

  Raises:
    ValueError:  If `states` and `filter_threshold` or `states` and
      `filter_beyond_lag` are both lists with different lengths.

  #### Examples

  We use ESS to estimate standard error.

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

  target = tfd.MultivariateNormalDiag(scale_diag=[1., 2.])

  # Get 1000 states from one chain.
  states = tfp.mcmc.sample_chain(
      num_burnin_steps=200,
      num_results=1000,
      current_state=tf.constant([0., 0.]),
      kernel=tfp.mcmc.HamiltonianMonteCarlo(
        target_log_prob_fn=target.log_prob,
        step_size=0.05,
        num_leapfrog_steps=20))
  states.shape
  ==> (1000, 2)

  ess = Func(states)
  ==> Shape (2,) Tensor

  mean, variance = tf.nn.moments(states, axis=0)
  standard_error = tf.sqrt(variance / ess)
  ```

  Some math shows that, with `R_k` the auto-correlation sequence,
  `R_k := Covariance{X_1, X_{1+k}} / Variance{X_1}`, we have

  ```ESS(N) =  N / [ 1 + 2 * ( (N - 1) / N * R_1 + ... + 1 / N * R_{N-1}  ) ]```

  This function estimates the above by first estimating the auto-correlation.
  Since `R_k` must be estimated using only `N - k` samples, it becomes
  progressively noisier for larger `k`.  For this reason, the summation over
  `R_k` should be truncated at some number `filter_beyond_lag < N`.  Since many
  MCMC methods generate chains where `R_k > 0`, a reasonable criteria is to
  truncate at the first index where the estimated auto-correlation becomes
  negative.

  The arguments `filter_beyond_lag`, `filter_threshold` are filters intended to
  remove noisy tail terms from `R_k`.  They combine in an "OR" manner meaning
  terms are removed if they were to be filtered under the `filter_beyond_lag` OR
  `filter_threshold` criteria.
  """
  arg_4 = _is_list_like(arg_0)

  # Convert all args to lists.
  if not arg_4:
    arg_0 = [arg_0]

  arg_2 = _broadcast_maybelist_arg(arg_0, arg_2,
                                               'filter_beyond_lag')
  arg_1 = _broadcast_maybelist_arg(arg_0, arg_1,
                                              'filter_threshold')

  # Process items, one at a time.
  with tf.compat.v1.name_scope(arg_3, 'Func'):
    arg_5 = [
        _Func_single_state(s, ml, mlt)
        for (s, ml, mlt) in zip(arg_0, arg_2, arg_1)
    ]

  if arg_4:
    return arg_5
  return arg_5[0]


def Func(arg_0, arg_1,
                                        arg_2):
  """ESS computation for one single Tensor argument."""

  with tf.compat.v1.name_scope(
      'effective_sample_size_single_state',
      values=[arg_0, arg_1, arg_2]):

    arg_0 = tf.convert_to_tensor(value=arg_0, name='states')
    arg_3 = arg_0.dtype

    # filter_beyond_lag == None ==> auto_corr is the full sequence.
    arg_4 = stats.auto_correlation(
        arg_0, axis=0, max_lags=arg_1)
    if arg_2 is not None:
      arg_2 = tf.convert_to_tensor(
          value=arg_2, dtype=arg_3, name='filter_threshold')
      # Get a binary mask to zero out values of auto_corr below the threshold.
      #   mask[i, ...] = 1 if auto_corr[j, ...] > threshold for all j <= i,
      #   mask[i, ...] = 0, otherwise.
      # So, along dimension zero, the mask will look like [1, 1, ..., 0, 0,...]
      # Building step by step,
      #   Assume auto_corr = [1, 0.5, 0.0, 0.3], and filter_threshold = 0.2.
      # Step 1:  mask = [False, False, True, False]
      arg_5 = arg_4 < arg_2
      # Step 2:  mask = [0, 0, 1, 1]
      arg_5 = tf.cast(arg_5, dtype=arg_3)
      # Step 3:  mask = [0, 0, 1, 2]
      arg_5 = tf.cumsum(arg_5, axis=0)
      # Step 4:  mask = [1, 1, 0, 0]
      arg_5 = tf.maximum(1. - arg_5, 0.)
      arg_4 *= arg_5

    # With R[k] := auto_corr[k, ...],
    # ESS = N / {1 + 2 * Sum_{k=1}^N (N - k) / N * R[k]}
    #     = N / {-1 + 2 * Sum_{k=0}^N (N - k) / N * R[k]} (since R[0] = 1)
    #     approx N / {-1 + 2 * Sum_{k=0}^M (N - k) / N * R[k]}
    # where M is the filter_beyond_lag truncation point chosen above.

    # Get the factor (N - k) / N, and give it shape [M, 1,...,1], having total
    # ndims the same as auto_corr
    arg_6 = _axis_size(arg_0, axis=0)
    arg_7 = tf.range(0., _axis_size(arg_4, axis=0))
    arg_8 = (arg_6 - arg_7) / arg_6
    if arg_4.shape.ndims is not None:
      arg_9 = [-1] + [1] * (arg_4.shape.ndims - 1)
    else:
      arg_9 = tf.concat(
          ([-1],
           tf.ones([tf.rank(arg_4) - 1], dtype=tf.int32)),
          axis=0)
    arg_8 = tf.reshape(arg_8, arg_9)

    return arg_6 / (-1 +
                2 * tf.reduce_sum(input_tensor=arg_8 * arg_4, axis=0))


def Func(arg_0, arg_1):
  """potential_scale_reduction for one single state `Tensor`."""
  with tf.compat.v1.name_scope(
      'potential_scale_reduction_single_state',
      values=[arg_0, arg_1]):
    # We assume exactly one leading dimension indexes e.g. correlated samples
    # from each Markov chain.
    arg_0 = tf.convert_to_tensor(value=arg_0, name='state')
    arg_2 = 1

    arg_3 = tf.range(0, arg_2)
    arg_4 = tf.range(arg_2,
                          arg_2 + arg_1)
    arg_5 = tf.range(
        0, arg_2 + arg_1)

    arg_6 = _axis_size(arg_0, arg_3)
    arg_7 = _axis_size(arg_0, arg_4)

    # In the language of Brooks and Gelman (1998),
    # B / n is the between chain variance, the variance of the chain means.
    # W is the within sequence variance, the mean of the chain variances.
    arg_8 = _reduce_variance(
        tf.reduce_mean(input_tensor=arg_0, axis=arg_3, keepdims=True),
        arg_5,
        biased=False)
    arg_9 = tf.reduce_mean(
        input_tensor=_reduce_variance(
            arg_0, arg_3, keepdims=True, biased=True),
        axis=arg_5)

    # sigma^2_+ is an estimate of the true variance, which would be unbiased if
    # each chain was drawn from the target.  c.f. "law of total variance."
    arg_10 = arg_9 + arg_8

    return ((arg_7 + 1.) / arg_7) * arg_10 / arg_9 - (arg_6 - 1.) / (arg_7 * arg_6)


def Func(arg_0, arg_1, arg_2):
  """Broadcast a listable secondary_arg to that of states."""
  if _is_list_like(arg_1):
    if len(arg_1) != len(arg_0):
      raise ValueError('Argument `%s` was a list of different length ({}) than '
                       '`states` ({})'.format(arg_2, len(arg_0)))
  else:
    arg_1 = [arg_1] * len(arg_0)

  return arg_1


def Func(
    arg_0, arg_1, arg_2,
    arg_3=False, arg_4=None):
  """Use LogNormal quantiles to form quadrature on positive-reals.

  Args:
    loc: `float`-like (batch of) scalar `Tensor`; the location parameter of
      the LogNormal prior.
    scale: `float`-like (batch of) scalar `Tensor`; the scale parameter of
      the LogNormal prior.
    quadrature_size: Python `int` scalar representing the number of quadrature
      points.
    validate_args: Python `bool`, default `False`. When `True` distribution
      parameters are checked for validity despite possibly degrading runtime
      performance. When `False` invalid inputs may silently render incorrect
      outputs.
    name: Python `str` name prefixed to Ops created by this class.

  Returns:
    grid: (Batch of) length-`quadrature_size` vectors representing the
      `log_rate` parameters of a `Poisson`.
    probs: (Batch of) length-`quadrature_size` vectors representing the
      weight associate with each `grid` value.
  """
  with tf.name_scope(arg_4 or "Func"):
    # Create a LogNormal distribution.
    arg_5 = transformed_distribution.TransformedDistribution(
        distribution=normal.Normal(arg_0=arg_0, arg_1=arg_1),
        bijector=exp_bijector.Exp(),
        arg_3=arg_3)
    arg_6 = tensorshape_util.rank(arg_5.batch_shape)
    if arg_6 is None:
      arg_6 = tf.shape(input=arg_5.batch_shape_tensor())[0]

    def _compute_quantiles():
      """Helper to build quantiles."""
      # Omit {0, 1} since they might lead to Inf/NaN.
      arg_7 = tf.zeros([], dtype=arg_5.dtype)
      arg_8 = tf.linspace(arg_7, 1., arg_2 + 3)[1:-1]
      # Expand edges so its broadcast across batch dims.
      arg_8 = tf.reshape(
          arg_8,
          shape=tf.concat(
              [[-1], tf.ones([arg_6], dtype=tf.int32)], axis=0))
      arg_9 = arg_5.quantile(arg_8)
      # Cyclically permute left by one.
      arg_10 = tf.concat([tf.range(1, 1 + arg_6), [0]], axis=0)
      arg_9 = tf.transpose(a=arg_9, arg_10=arg_10)
      return arg_9
    arg_9 = _compute_quantiles()

    # Compute grid as quantile midpoints.
    arg_11 = (arg_9[..., :-1] + arg_9[..., 1:]) / 2.
    # Set shape hints.
    arg_12 = tensorshape_util.concatenate(arg_5.batch_shape,
                                             [arg_2])
    tensorshape_util.set_shape(arg_11, arg_12)

    # By construction probs is constant, i.e., `1 / quadrature_size`. This is
    # important, because non-constant probs leads to non-reparameterizable
    # samples.
    arg_13 = tf.fill(
        dims=[arg_2], value=1. / tf.cast(arg_2, arg_5.dtype))

    return arg_11, arg_13


def Func(arg_0, arg_1, arg_2, arg_3=False):
    """Helper to merge which handles merging one value."""
    if arg_1 is None:
      return arg_2
    if arg_2 is None:
      return arg_1
    if (arg_1 == arg_2) if arg_3 else (arg_1 is arg_2):
      return arg_1
    raise ValueError("Incompatible values: %s != %s" % (arg_1, arg_2))


def Func(arg_0, arg_1):
    """Converts nested `tuple`, `list`, or `dict` to nested `tuple`."""
    if isinstance(arg_1, dict):
      return arg_0.Func(tuple(sorted(arg_1.items())))
    elif isinstance(arg_1, (list, tuple)):
      return tuple(map(arg_0.Func, arg_1))

    return arg_1


def Func(arg_0, arg_1, arg_2, arg_3=None,
                              arg_4=None):
  """Computes the doubling increments for the left end point.

  The doubling procedure expands an initial interval to find a superset of the
  true slice. At each doubling iteration, the interval width is doubled to
  either the left or the right hand side with equal probability.
  If, initially, the left end point is at `L(0)` and the width of the
  interval is `w(0)`, then the left end point and the width at the
  k-th iteration (denoted L(k) and w(k) respectively) are given by the following
  recursions:

  ```none
  w(k) = 2 * w(k-1)
  L(k) = L(k-1) - w(k-1) * X_k, X_k ~ Bernoulli(0.5)
  or, L(0) - L(k) = w(0) Sum(2^i * X(i+1), 0 <= i < k)
  ```

  This function computes the sequence of `L(0)-L(k)` and `w(k)` for k between 0
  and `max_doublings` independently for each chain.

  Args:
    batch_shape: Positive int32 `tf.Tensor`. The batch shape.
    max_doublings: Scalar positive int32 `tf.Tensor`. The maximum number of
      doublings to consider.
    step_size: A real `tf.Tensor` with shape compatible with [num_chains].
      The size of the initial interval.
    seed: (Optional) positive int. The random seed. If None, no seed is set.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., 'find_slice_bounds').

  Returns:
    left_increments: A tensor of shape (max_doublings+1, batch_shape). The
      relative position of the left end point after the doublings.
    widths: A tensor of shape (max_doublings+1, ones_like(batch_shape)). The
      widths of the intervals at each stage of the doubling.
  """
  with tf.compat.v1.name_scope(arg_4, 'left_doubling_increments',
                               [arg_0, arg_1, arg_2]):

    arg_2 = tf.convert_to_tensor(value=arg_2)
    arg_5 = arg_2.dtype.base_dtype
    # Output shape of the left increments tensor.
    arg_6 = tf.concat(([arg_1 + 1], arg_0), axis=0)
    # A sample realization of X_k.
    arg_7 = distributions.Bernoulli(0.5, arg_5=arg_5).sample(
        sample_shape=arg_6, arg_3=arg_3)

    # The widths of the successive intervals. Starts with 1.0 and ends with
    # 2^max_doublings.
    arg_8 = tf.cast(2 ** tf.range(0, arg_1+1), arg_5=arg_5)
    # Output shape of the `widths` tensor.
    arg_9 = tf.concat(([arg_1 + 1],
                              tf.ones_like(arg_0)), axis=0)
    arg_8 = tf.reshape(arg_8, shape=arg_9)
    # Widths shape is [max_doublings + 1, 1, 1, 1...].
    arg_10 = arg_8 * arg_2

    # Take the cumulative sum of the left side increments in slice width to give
    # the resulting distance from the inital lower bound.
    arg_11 = tf.cumsum(arg_10 * arg_7, exclusive=True, axis=0)
    return arg_11, arg_10


def Func(arg_0, arg_1=None):
  """Finds the index of the optimal set of bounds for each chain.

  For each chain, finds the smallest set of bounds for which both edges lie
  outside the slice. This is equivalent to the point at which a for loop
  implementation (P715 of Neal (2003)) of the algorithm would terminate.

  Performs the following calculation, where i is the number of doublings that
  have been performed and k is the max number of doublings:

  (2 * k - i) * flag + i

  The argmax of the above returns the earliest index where the bounds were
  outside the slice and if there is no such point, the widest bounds.

  Args:
    x: A tensor of shape (max_doublings+1, batch_shape). Type int32, with value
      0 or 1. Indicates if this set of bounds is outside the slice.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., 'find_slice_bounds').

  Returns:
    indices: A tensor of shape batch_shape. Type int32, with the index of the
      first set of bounds outside the slice and if there are none, the index of
      the widest set.
  """
  with tf.compat.v1.name_scope(arg_1, 'find_best_interval_idx', [arg_0]):
    # Returns max_doublings + 1. Positive int32.
    arg_2 = tf.shape(input=arg_0)[0]
    arg_3 = arg_0.dtype.base_dtype
    # Factors by which to multiply the flag. Corresponds to (2 * k - i) above.
    arg_4 = tf.range(2 * arg_2, arg_2, -1, arg_3=arg_3)[:, tf.newaxis]
    # Factors by which to shift the flag. Corresponds to i above. Ensures the
    # widest bounds are selected if there are no bounds outside the slice.
    arg_5 = tf.range(arg_2, arg_3=arg_3)[:, tf.newaxis]
    arg_6 = tf.argmax(input=arg_4 * arg_0 + arg_5, axis=0, output_type=arg_3)
    return arg_6


def Func(arg_0,
                             arg_1,
                             arg_2,
                             arg_3,
                             arg_4,
                             arg_5=None,
                             arg_6=None):
  """Returns the bounds of the slice at each stage of doubling procedure.

  Precomputes the x coordinates of the left (L) and right (R) endpoints of the
  interval `I` produced in the "doubling" algorithm [Neal 2003][1] P713. Note
  that we simultaneously compute all possible doubling values for each chain,
  for the reason that at small-medium densities, the gains from parallel
  evaluation might cause a speed-up, but this will be benchmarked against the
  while loop implementation.

  Args:
    x_initial: `tf.Tensor` of any shape and any real dtype consumable by
      `target_log_prob`. The initial points.
    target_log_prob: A callable taking a `tf.Tensor` of shape and dtype as
      `x_initial` and returning a tensor of the same shape. The log density of
      the target distribution.
    log_slice_heights: `tf.Tensor` with the same shape as `x_initial` and the
      same dtype as returned by `target_log_prob`. The log of the height of the
      slice for each chain. The values must be bounded above by
      `target_log_prob(x_initial)`.
    max_doublings: Scalar positive int32 `tf.Tensor`. The maximum number of
      doublings to consider.
    step_size: `tf.Tensor` with same dtype as and shape compatible with
      `x_initial`. The size of the initial interval.
    seed: (Optional) positive int. The random seed. If None, no seed is set.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., 'find_slice_bounds').

  Returns:
    upper_bounds: A tensor of same shape and dtype as `x_initial`. Slice upper
      bounds for each chain.
    lower_bounds: A tensor of same shape and dtype as `x_initial`. Slice lower
      bounds for each chain.
    both_ok: A tensor of shape `x_initial` and boolean dtype. Indicates if both
      the chosen upper and lower bound lie outside of the slice.

  #### References

  [1]: Radford M. Neal. Slice Sampling. The Annals of Statistics. 2003, Vol 31,
       No. 3 , 705-767.
       https://projecteuclid.org/download/pdf_1/euclid.aos/1056562461
  """
  with tf.compat.v1.name_scope(
      arg_6, 'Func',
      [arg_0, arg_2, arg_3, arg_4]):
    arg_7 = distributions.SeedStream(arg_5, salt='Func')
    arg_0 = tf.convert_to_tensor(value=arg_0)
    arg_8 = tf.shape(input=arg_0)
    arg_9 = arg_4.dtype.base_dtype
    arg_10 = arg_0 + arg_4 * tf.random.uniform(
        arg_8, minval=-1.0, maxval=0.0, arg_9=arg_9, arg_5=arg_7())

    # Compute the increments by which we need to step the upper and lower bounds
    # part of the doubling procedure.
    arg_11, arg_12 = _left_doubling_increments(
        arg_8, arg_3, arg_4, arg_5=arg_7())
    # The left and right end points. Shape (max_doublings+1,) + batch_shape.
    arg_10 -= arg_11
    arg_13 = arg_10 + arg_12

    # Test if these end points lie outside of the slice.
    # Checks if the end points of the slice are outside the graph of the pdf.
    arg_14 = tf.map_fn(arg_1, arg_10)
    arg_15 = tf.map_fn(arg_1, arg_13)
    arg_16 = arg_14 < arg_2
    arg_17 = arg_15 < arg_2
    arg_18 = arg_16 & arg_17

    arg_19 = tf.reshape(arg_18, [arg_3 + 1, -1])

    arg_20 = _find_best_interval_idx(
        tf.cast(arg_19, arg_9=tf.int32))

    # Formats the above index as required to use with gather_nd.
    arg_21 = tf.stack(
        [arg_20,
         tf.range(tf.size(input=arg_20))],
        axis=1,
        arg_6='point_index_gather')
    arg_22 = tf.reshape(arg_10, [arg_3 + 1, -1])
    arg_23 = tf.reshape(arg_13, [arg_3 + 1, -1])
    # The x values of the uppper and lower bounds of the slices for each chain.
    arg_24 = tf.reshape(tf.gather_nd(arg_22, arg_21),
                              arg_8)
    arg_25 = tf.reshape(tf.gather_nd(arg_23, arg_21),
                              arg_8)
    arg_18 = tf.reduce_any(input_tensor=arg_18, axis=0)
    return arg_25, arg_24, arg_18


def Func(arg_0, arg_1, arg_2,
                           arg_3, arg_4, arg_5, arg_6=None,
                           arg_7=None):
  """Samples from the slice by applying shrinkage for rejected points.

  Implements the one dimensional slice sampling algorithm of Neal (2003), with a
  doubling algorithm (Neal 2003 P715 Fig. 4), which doubles the size of the
  interval at each iteration and shrinkage (Neal 2003 P716 Fig. 5), which
  reduces the width of the slice when a selected point is rejected, by setting
  the relevant bound that that value. Randomly sampled points are checked for
  two criteria: that they lie within the slice and that they pass the
  acceptability check (Neal 2003 P717 Fig. 6), which tests that the new state
  could have generated the previous one.

  Args:
    x_initial: A tensor of any shape. The initial positions of the chains. This
      function assumes that all the dimensions of `x_initial` are batch
      dimensions (i.e. the event shape is `[]`).
    target_log_prob: Callable accepting a tensor like `x_initial` and returning
      a tensor containing the log density at that point of the same shape.
    log_slice_heights: Tensor of the same shape and dtype as the return value
      of `target_log_prob` when applied to `x_initial`. The log of the height of
      the chosen slice.
    step_size: A tensor of shape and dtype compatible with `x_initial`. The min
      interval size in the doubling algorithm.
    lower_bounds: Tensor of same shape and dtype as `x_initial`. Slice lower
      bounds for each chain.
    upper_bounds: Tensor of same shape and dtype as `x_initial`. Slice upper
      bounds for each chain.
    seed: (Optional) positive int. The random seed. If None, no seed is set.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., 'find_slice_bounds').

  Returns:
    x_proposed: A tensor of the same shape and dtype as `x_initial`. The next
      proposed state of the chain.
  """
  with tf.compat.v1.name_scope(
      arg_7, 'sample_with_shrinkage',
      [arg_0, arg_2, arg_3, arg_4, arg_5]):
    arg_8 = distributions.SeedStream(arg_6, salt='Func')
    # Keeps track of whether an acceptable sample has been found for the chain.
    arg_9 = tf.zeros_like(arg_0, dtype=tf.bool)
    arg_10 = lambda arg_9, *ignored_args: ~tf.reduce_all(input_tensor=arg_9)
    arg_11 = tf.identity(arg_0)
    arg_12 = tf.shape(input=arg_0)
    arg_13 = arg_0.dtype.base_dtype
    def _body(arg_9, arg_14, arg_15, arg_11):
      """Iterates until every chain has found a suitable next state."""
      arg_16 = tf.random.uniform(
          arg_12, dtype=arg_13, arg_6=arg_8())
      arg_17 = tf.where(~arg_9, arg_14 + arg_16 * (arg_15 - arg_14), arg_11)
      arg_18 = _test_acceptance(arg_0, arg_1=arg_1,
                                    decided=arg_9,
                                    arg_2=arg_2,
                                    arg_17=arg_17, arg_3=arg_3,
                                    arg_4=arg_14, arg_5=arg_15)
      arg_19 = arg_2 < arg_1(arg_17)
      arg_20 = arg_19 & arg_18
      arg_21 = arg_9 | arg_20
      # Note that it might seem that we are moving the left and right end points
      # even if the point has been accepted (which is contrary to the stated
      # algorithm in Neal). However, this does not matter because the endpoints
      # for points that have been already accepted are not used again so it
      # doesn't matter what we do with them.
      arg_22 = tf.where(arg_17 < arg_0, arg_17, arg_14)
      arg_23 = tf.where(arg_17 >= arg_0, arg_17, arg_15)
      return arg_21, arg_22, arg_23, arg_17

    return tf.while_loop(
        arg_10=arg_10,
        body=_body,
        loop_vars=(arg_9, arg_4, arg_5, arg_11))[-1]