def Func(**arg_0):
  """Decorator to programmatically expand the docstring.

  Args:
    **kwargs: Keyword arguments to set. For each key-value pair `k` and `v`,
      the key is found as `${k}` in the docstring and replaced with `v`.

  Returns:
    Decorated function.
  """
  def _fn_wrapped(arg_1):
    """Original function with modified `__doc__` attribute."""
    arg_2 = inspect.cleandoc(arg_1.__doc__)
    for arg_3, arg_4 in six.iteritems(arg_0):
      # Capture each ${k} reference to replace with v.
      # We wrap the replacement in a function so no backslash escapes
      # are processed.
      arg_5 = r'\$\{' + str(arg_3) + r'\}'
      arg_2 = re.sub(arg_5, lambda match: arg_4, arg_2)  # pylint: disable=cell-var-from-loop
    arg_1.__doc__ = arg_2
    return arg_1
  return _fn_wrapped


def Func(arg_0):
  """Infer the original name passed into a distribution constructor.

  Distributions typically follow the pattern of
  with.name_scope(name) as name:
    super(name=name)
  so we attempt to reverse the name-scope transformation to allow
  addressing of RVs by the distribution's original, user-visible
  name kwarg.

  Args:
    distribution: a tfd.Distribution instance.
  Returns:
    simple_name: the original name passed into the Distribution.

  #### Example

  ```
  d1 = tfd.Normal(0., 1., name='x') # d1.name = 'x/'
  d2 = tfd.Normal(0., 1., name='x') # d2.name = 'x_2/'
  Func(d2) # returns 'x'

  ```

  """
  arg_1 = arg_0.name

  # turn 'scope/x/' into 'x'
  if arg_1.endswith('/'):
    arg_1 = arg_1.split('/')[-2]

  # turn 'x_3' into 'x'
  arg_2 = arg_1.split('_')
  if arg_2[-1].isdigit():
    arg_1 = '_'.join(arg_2[:-1])

  return arg_1


def Func(arg_0, arg_1, arg_2, arg_3):
  """RandomVariable constructor with a dummy name argument."""
  # Program transformations (e.g., `make_log_joint_fn`) assume that
  # the traced constructor has `name` and `value` kwargs, enabling
  # them to override the value of an RV according to its name.
  # User-defined RVs inherit their name from the provided
  # distribution; this helper method exposes the name as a dummy kwarg
  # so that it's visible to program transformations.
  del arg_3  # unused
  return RandomVariable(arg_0=arg_0,
                        arg_1=arg_1,
                        arg_2=arg_2)


def Func(arg_0):
  """Factory function to make random variable given distribution class."""

  @interceptable
  @functools.wraps(arg_0, assigned=('__module__', '__name__'))
  @docstring_util.expand_docstring(
      cls=arg_0.__name__,
      doc=inspect.cleandoc(arg_0.__init__.__doc__ or ''))
  def func(*arg_1, **arg_2):
    # pylint: disable=g-doc-args
    """Create a random variable for ${cls}.

    See ${cls} for more details.

    Returns:
      RandomVariable.

    #### Original Docstring for Distribution

    ${doc}
    """
    # pylint: enable=g-doc-args
    arg_3 = arg_2.pop('sample_shape', ())
    arg_4 = arg_2.pop('value', None)
    return RandomVariable(distribution=arg_0(*arg_1, **arg_2),
                          arg_3=arg_3,
                          arg_4=arg_4)
  return func


def Func(arg_0, arg_1, arg_2):
  """Compute one-step-ahead predictive distributions for all timesteps.

  Given samples from the posterior over parameters, return the predictive
  distribution over observations at each time `T`, given observations up
  through time `T-1`.

  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` including a
      mask `Tensor` to encode the locations of missing observations.
    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:
    forecast_dist: a `tfd.MixtureSameFamily` instance with event shape
      [num_timesteps] and
      batch shape `concat([sample_shape, model.batch_shape])`, with
      `num_posterior_draws` mixture components. The `t`th step represents the
      forecast distribution `p(observed_time_series[t] |
      observed_time_series[0:t-1], parameter_samples)`.

  #### Examples

  Suppose we've built a model and fit it to data using HMC:

  ```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)

    samples, kernel_results = tfp.sts.fit_with_hmc(model, observed_time_series)
  ```

  Passing the posterior samples into `Func`, we construct a
  one-step-ahead predictive distribution:

  ```python
    Func_dist = tfp.sts.Func(
      model, observed_time_series, parameter_samples=samples)

    predictive_means = Func_dist.mean()
    predictive_scales = Func_dist.stddev()
  ```

  If using variational inference instead of HMC, we'd construct a forecast using
  samples from the variational posterior:

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

    # OMITTED: take steps to optimize variational loss

    samples = {k: q.sample(30) for (k, q) in variational_distributions.items()}
    Func_dist = tfp.sts.Func(
      model, observed_time_series, parameter_samples=samples)
  ```

  We can visualize the forecast by plotting:

  ```python
    from matplotlib import pylab as plt
    def plot_Func(observed_time_series,
                                 forecast_mean,
                                 forecast_scale):
      plt.figure(figsize=(12, 6))
      num_timesteps = forecast_mean.shape[-1]
      c1, c2 = (0.12, 0.47, 0.71), (1.0, 0.5, 0.05)
      plt.plot(observed_time_series, label="observed time series", color=c1)
      plt.plot(forecast_mean, label="one-step prediction", color=c2)
      plt.fill_between(np.arange(num_timesteps),
                       forecast_mean - 2 * forecast_scale,
                       forecast_mean + 2 * forecast_scale,
                       alpha=0.1, color=c2)
      plt.legend()

    plot_Func(observed_time_series,
                             forecast_mean=predictive_means,
                             forecast_scale=predictive_scales)
  ```

  To detect anomalous timesteps, we check whether the observed value at each
  step is within a 95% predictive interval, i.e., two standard deviations from
  the mean:

  ```python
    z_scores = ((observed_time_series[..., 1:] - predictive_means[..., :-1])
                 / predictive_scales[..., :-1])
    anomalous_timesteps = tf.boolean_mask(
        tf.range(1, num_timesteps),
        tf.abs(z_scores) > 2.0)
  ```

  """

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

    [
        arg_1,
        arg_3
    ] = sts_util.canonicalize_observed_time_series_with_mask(
        arg_1)

    # Run filtering over the training timesteps to extract the
    # predictive means and variances.
    arg_4 = dist_util.prefer_static_value(
        tf.shape(input=arg_1))[-2]
    arg_5 = arg_0.make_state_space_model(
        arg_4=arg_4, param_vals=arg_2)
    (arg_6, arg_6, arg_6, arg_6, arg_6, arg_7, arg_8
    ) = arg_5.forward_filter(arg_1, mask=arg_3)

    # Squeeze dims to convert from LGSSM's event shape `[num_timesteps, 1]`
    # to a scalar time series.
    return sts_util.mix_over_posterior_draws(
        means=arg_7[..., 0],
        variances=arg_8[..., 0, 0])


def Func(arg_0,
             arg_1,
             arg_2,
             arg_3):
  """Construct predictive distribution over future observations.

  Given samples from the posterior over parameters, return the predictive
  distribution over future observations for num_steps_Func timesteps.

  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` including a
      mask `Tensor` to encode the locations of missing observations.
    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.
    num_steps_Func: scalar `int` `Tensor` number of steps to Func.

  Returns:
    Func_dist: a `tfd.MixtureSameFamily` instance with event shape
      [num_steps_Func, 1] and batch shape
      `concat([sample_shape, model.batch_shape])`, with `num_posterior_draws`
      mixture components.

  #### Examples

  Suppose we've built a model and fit it to data using HMC:

  ```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)

    samples, kernel_results = tfp.sts.fit_with_hmc(model, observed_time_series)
  ```

  Passing the posterior samples into `Func`, we construct a Func
  distribution:

  ```python
    Func_dist = tfp.sts.Func(model, observed_time_series,
                                     parameter_samples=samples,
                                     num_steps_Func=50)

    Func_mean = Func_dist.mean()[..., 0]  # shape: [50]
    Func_scale = Func_dist.stddev()[..., 0]  # shape: [50]
    Func_samples = Func_dist.sample(10)[..., 0]  # shape: [10, 50]
  ```

  If using variational inference instead of HMC, we'd construct a Func using
  samples from the variational posterior:

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

    # OMITTED: take steps to optimize variational loss

    samples = {k: q.sample(30) for (k, q) in variational_distributions.items()}
    Func_dist = tfp.sts.Func(model, observed_time_series,
                                         parameter_samples=samples,
                                         num_steps_Func=50)
  ```

  We can visualize the Func by plotting:

  ```python
    from matplotlib import pylab as plt
    def plot_Func(observed_time_series,
                      Func_mean,
                      Func_scale,
                      Func_samples):
      plt.figure(figsize=(12, 6))

      num_steps = observed_time_series.shape[-1]
      num_steps_Func = Func_mean.shape[-1]
      num_steps_train = num_steps - num_steps_Func

      c1, c2 = (0.12, 0.47, 0.71), (1.0, 0.5, 0.05)
      plt.plot(np.arange(num_steps), observed_time_series,
               lw=2, color=c1, label='ground truth')

      Func_steps = np.arange(num_steps_train,
                       num_steps_train+num_steps_Func)
      plt.plot(Func_steps, Func_samples.T, lw=1, color=c2, alpha=0.1)
      plt.plot(Func_steps, Func_mean, lw=2, ls='--', color=c2,
               label='Func')
      plt.fill_between(Func_steps,
                       Func_mean - 2 * Func_scale,
                       Func_mean + 2 * Func_scale, color=c2, alpha=0.2)

      plt.xlim([0, num_steps])
      plt.legend()

    plot_Func(observed_time_series,
                  Func_mean=Func_mean,
                  Func_scale=Func_scale,
                  Func_samples=Func_samples)
  ```

  """

  with tf.compat.v1.name_scope(
      'Func',
      values=[arg_1, arg_2, arg_3]):
    [
        arg_1,
        arg_4
    ] = sts_util.canonicalize_observed_time_series_with_mask(
        arg_1)

    # Run filtering over the observed timesteps to extract the
    # latent state posterior at timestep T+1 (i.e., the final
    # filtering distribution, pushed through the transition model).
    # This is the prior for the Func model ("today's prior
    # is yesterday's posterior").
    arg_5 = dist_util.prefer_static_value(
        tf.shape(input=arg_1))[-2]
    arg_6 = arg_0.make_state_space_model(
        num_timesteps=arg_5, param_vals=arg_2)
    (arg_7, arg_7, arg_7, arg_8, arg_9, arg_7, arg_7
    ) = arg_6.forward_filter(arg_1, arg_4=arg_4)

    # Build a batch of state-space models over the Func period. Because
    # we'll use MixtureSameFamily to mix over the posterior draws, we need to
    # do some shenanigans to move the `[num_posterior_draws]` batch dimension
    # from the leftmost to the rightmost side of the model's batch shape.
    # TODO(b/120245392): enhance `MixtureSameFamily` to reduce along an
    # arbitrary axis, and eliminate `move_dimension` calls here.
    arg_2 = arg_0._canonicalize_param_vals_as_map(arg_2)  # pylint: disable=protected-access
    arg_10 = {
        param.name: dist_util.move_dimension(
            arg_2[param.name],
            0, -(1 + _prefer_static_event_ndims(param.prior)))
        for param in arg_0.parameters}
    arg_11 = tfd.MultivariateNormalFullCovariance(
        loc=dist_util.move_dimension(arg_8[..., -1, :], 0, -2),
        covariance_matrix=dist_util.move_dimension(
            arg_9[..., -1, :, :], 0, -3))

    # Ugly hack: because we moved `num_posterior_draws` to the trailing (rather
    # than leading) dimension of parameters, the parameter batch shapes no
    # longer broadcast against the `constant_offset` attribute used in `sts.Sum`
    # models. We fix this by manually adding an extra broadcasting dim to
    # `constant_offset` if present.
    # The root cause of this hack is that we mucked with param dimensions above
    # and are now passing params that are 'invalid' in the sense that they don't
    # match the shapes of the model's param priors. The fix (as above) will be
    # to update MixtureSameFamily so we can avoid changing param dimensions
    # altogether.
    # TODO(b/120245392): enhance `MixtureSameFamily` to reduce along an
    # arbitrary axis, and eliminate this hack.
    arg_12 = {}
    if hasattr(arg_0, 'constant_offset'):
      arg_12['constant_offset'] = tf.convert_to_tensor(
          value=arg_0.constant_offset,
          dtype=arg_11.dtype)[..., tf.newaxis]

    # We assume that any STS model that has a `constant_offset` attribute
    # will allow it to be overridden as a kwarg. This is currently just
    # `sts.Sum`.
    # TODO(b/120245392): when kwargs hack is removed, switch back to calling
    # the public version of `_make_state_space_model`.
    arg_13 = arg_0._make_state_space_model(  # pylint: disable=protected-access
        num_timesteps=arg_3,
        param_map=arg_10,
        initial_state_prior=arg_11,
        initial_step=arg_5,
        **arg_12)

    arg_14 = dist_util.prefer_static_value(
        arg_13.batch_shape_tensor())[-1]
    return tfd.MixtureSameFamily(
        mixture_distribution=tfd.Categorical(
            logits=tf.zeros([arg_14], dtype=arg_13.dtype)),
        components_distribution=arg_13)


def Func(arg_0, arg_1=-1, arg_2=False, arg_3=0):
  """Returns `max` or `mask` if `max` is not finite."""
  arg_4 = np.max(arg_0, arg_1=_astuple(arg_1), arg_2=arg_2)
  arg_5 = ~np.isfinite(arg_4)
  if arg_5.ndim > 0:
    arg_4[arg_5] = arg_3
  elif arg_5:
    arg_4 = arg_3
  return arg_4


def Func(arg_0, arg_1, arg_2=None, arg_3=None, arg_4=None,
                        arg_5=None):
  """Assert `x` has rank equal to `rank` or smaller.

  Example of adding a dependency to an operation:

  ```python
  with tf.control_dependencies([tf.Func(x, 2)]):
    output = tf.reduce_sum(x)
  ```

  Args:
    x:  Numeric `Tensor`.
    rank:  Scalar `Tensor`.
    data:  The tensors to print out if the condition is False.  Defaults to
      error message and first few entries of `x`.
    summarize: Print this many entries of each tensor.
    message: A string to prefix to the default message.
    name: A name for this operation (optional).
      Defaults to "Func".

  Returns:
    Op raising `InvalidArgumentError` unless `x` has specified rank or lower.
    If static checks determine `x` has correct rank, a `no_op` is returned.

  Raises:
    ValueError:  If static checks determine `x` has wrong rank.
  """
  with tf.compat.v2.name_scope(arg_5 or 'Func'):
    return tf.compat.v1.assert_less_equal(
        tf.rank(arg_0), arg_1, arg_2=arg_2, arg_3=arg_3, arg_4=arg_4)


def Func(arg_0, arg_1, arg_2=None):
  """OneHotCategorical helper computing probs, cdf, etc over its support."""
  with tf.compat.v1.name_scope(arg_2, 'eval_all_one_hot'):
    arg_3 = arg_1.event_shape_tensor()[-1]
    arg_4 = tf.size(input=arg_1.batch_shape_tensor())
    # Reshape `eye(d)` to: `[d] + [1]*batch_ndims + [d]`.
    arg_5 = tf.reshape(
        tf.eye(arg_3, dtype=arg_1.dtype),
        shape=tf.pad(
            tensor=tf.ones(arg_4, tf.int32),
            paddings=[[1, 1]],
            constant_values=arg_3))
    # Compute `fn(x)` then cyclically left-transpose one dim.
    arg_6 = tf.pad(tensor=tf.range(1, arg_4 + 1), paddings=[[0, 1]])
    return tf.transpose(a=arg_0(arg_1, arg_5), arg_6=arg_6)


def Func(arg_0):
  """Return a convert-to-tensor func, given a name, config, callable, etc."""
  if arg_0 is None:
    return None

  if isinstance(arg_0, six.string_types):
    arg_0 = str(arg_0)
    return _deserialize(arg_0)

  if isinstance(arg_0, dict):
    return _deserialize(arg_0)

  if isinstance(arg_0, property):
    arg_0 = arg_0.fget
  if callable(arg_0):
    return arg_0

  raise ValueError('Could not interpret '
                   'convert-to-tensor function identifier:', arg_0)


def Func():
  """Yields the top-most interceptor on the thread-local interceptor stack.

  Operations may be intercepted by multiple nested interceptors. Once reached,
  an operation can be forwarded through nested interceptors until resolved.
  To allow for nesting, implement interceptors by re-wrapping their first
  argument (`f`) as an `interceptable`. To avoid nesting, manipulate the
  computation without using `interceptable`.

  This function allows for nesting by manipulating the thread-local interceptor
  stack, so that operations are intercepted in the order of interceptor nesting.

  #### Examples

  ```python
  from tensorflow_probability import edward2 as ed

  def model():
    x = ed.Normal(loc=0., scale=1., name="x")
    y = ed.Normal(loc=x, scale=1., name="y")
    return x + y

  def double(f, *args, **kwargs):
    return 2. * interceptable(f)(*args, **kwargs)

  def set_y(f, *args, **kwargs):
    if kwargs.get("name") == "y":
      kwargs["value"] = 0.42
    return interceptable(f)(*args, **kwargs)

  with interception(double):
    with interception(set_y):
      z = model()
  ```

  This will firstly put `double` on the stack, and then `set_y`,
  resulting in the stack:
  (TOP) set_y -> double -> apply (BOTTOM)

  The execution of `model` is then (top lines are current stack state):
  1) (TOP) set_y -> double -> apply (BOTTOM);
  `ed.Normal(0., 1., "x")` is intercepted by `set_y`, and as the name is not "y"
  the operation is simply forwarded to the next interceptor on the stack.

  2) (TOP) double -> apply (BOTTOM);
  `ed.Normal(0., 1., "x")` is intercepted by `double`, to produce
  `2*ed.Normal(0., 1., "x")`, with the operation being forwarded down the stack.

  3) (TOP) apply (BOTTOM);
  `ed.Normal(0., 1., "x")` is intercepted by `apply`, which simply calls the
  constructor.

  (At this point, the nested calls to `Func()`, produced by
  forwarding operations, exit, and the current stack is again:
  (TOP) set_y -> double -> apply (BOTTOM))

  4) (TOP) set_y -> double -> apply (BOTTOM);
  `ed.Normal(0., 1., "y")` is intercepted by `set_y`,
  the value of `y` is set to 0.42 and the operation is forwarded down the stack.

  5) (TOP) double -> apply (BOTTOM);
  `ed.Normal(0., 1., "y")` is intercepted by `double`, to produce
  `2*ed.Normal(0., 1., "y")`, with the operation being forwarded down the stack.

  6) (TOP) apply (BOTTOM);
  `ed.Normal(0., 1., "y")` is intercepted by `apply`, which simply calls the
  constructor.

  The final values for `x` and `y` inside of `model()` are tensors where `x` is
  a random draw from Normal(0., 1.) doubled, and `y` is a constant 0.84, thus
  z = 2 * Normal(0., 1.) + 0.84.
  """
  try:
    arg_0 = _interceptor_stack.stack.pop()
    yield arg_0
  finally:
    _interceptor_stack.stack.append(arg_0)


def Func(arg_0):
  """Decorator that wraps `func` so that its execution is intercepted.

  The wrapper passes `func` to the interceptor for the current thread.

  If there is no next interceptor, we perform an "immediate" call to `func`.
  That is, `func` terminates without forwarding its execution to another
  interceptor.

  Args:
    func: Function to wrap.

  Returns:
    The decorated function.
  """
  @functools.wraps(arg_0)
  def func_wrapped(*arg_1, **arg_2):
    with get_next_interceptor() as interceptor:
      return interceptor(arg_0, *arg_1, **arg_2)

  return func_wrapped


def Func(arg_0, arg_1=2, arg_2=5.0):
  """Generates synthetic data for binary classification.

  Args:
    num_examples: The number of samples to generate (scalar Python `int`).
    input_size: The input space dimension (scalar Python `int`).
    weights_prior_stddev: The prior standard deviation of the weight
      vector. (scalar Python `float`).

  Returns:
    random_weights: Sampled weights as a Numpy `array` of shape
      `[input_size]`.
    random_bias: Sampled bias as a scalar Python `float`.
    design_matrix: Points sampled uniformly from the cube `[-1,
       1]^{input_size}`, as a Numpy `array` of shape `(num_examples,
       input_size)`.
    labels: Labels sampled from the logistic model `p(label=1) =
      logistic(dot(features, random_weights) + random_bias)`, as a Numpy
      `int32` `array` of shape `(num_examples, 1)`.
  """
  arg_3 = arg_2 * np.random.randn(arg_1)
  arg_4 = np.random.randn()
  arg_5 = np.random.rand(arg_0, arg_1) * 2 - 1
  arg_6 = np.reshape(
      np.dot(arg_5, arg_3) + arg_4,
      (-1, 1))
  arg_7 = 1. / (1 + np.exp(-arg_6))
  arg_8 = np.int32(arg_7 > np.random.rand(arg_0, 1))
  return arg_3, arg_4, np.float32(arg_5), arg_8


def Func(arg_0, arg_1, arg_2):
  """Build a Dataset iterator for supervised classification.

  Args:
    x: Numpy `array` of features, indexed by the first dimension.
    y: Numpy `array` of labels, with the same first dimension as `x`.
    batch_size: Number of elements in each training batch.

  Returns:
    batch_features: `Tensor` feed  features, of shape
      `[batch_size] + x.shape[1:]`.
    batch_labels: `Tensor` feed of labels, of shape
      `[batch_size] + y.shape[1:]`.
  """
  arg_3 = tf.data.Dataset.from_tensor_slices((arg_0, arg_1))
  arg_4 = arg_3.repeat().batch(arg_2)
  arg_5 = tf.compat.v1.data.make_one_shot_iterator(arg_4)
  arg_6, arg_7 = arg_5.get_next()
  return arg_6, arg_7


def Func(arg_0, arg_1):
  """Validate `map_values` if `validate_args`==True."""
  arg_2 = []

  arg_3 = 'Rank of map_values must be 1.'
  if tensorshape_util.rank(arg_0.shape) is not None:
    if tensorshape_util.rank(arg_0.shape) != 1:
      raise ValueError(arg_3)
  elif arg_1:
    arg_2.append(assert_util.assert_rank(arg_0, 1, arg_3=arg_3))

  arg_3 = 'Size of map_values must be greater than 0.'
  if tensorshape_util.num_elements(arg_0.shape) is not None:
    if tensorshape_util.num_elements(arg_0.shape) == 0:
      raise ValueError(arg_3)
  elif arg_1:
    arg_2.append(
        assert_util.assert_greater(
            tf.size(input=arg_0), 0, arg_3=arg_3))

  if arg_1:
    arg_2.append(
        assert_util.assert_equal(
            tf.math.is_strictly_increasing(arg_0),
            True,
            arg_3='map_values is not strictly increasing.'))

  return arg_2


def Func(arg_0: arg_1, arg_2: arg_3[arg_4[arg_5], arg_5]
                  ) -> arg_4[tf.Tensor, TensorNest, TensorNest]:
  """Calls `fn` and returns the gradients with respect to `fn`'s first output.

  Args:
    fn: A `TransitionOperator`.
    args: Arguments to `fn`

  Returns:
    ret: First output of `fn`.
    extra: Second output of `fn`.
    grads: Gradients of `ret` with respect to `args`.
  """
  with tf.GradientTape() as tape:
    tape.watch(arg_2)
    arg_6, arg_7 = call_fn(arg_0, arg_2)
  arg_8 = tape.gradient(arg_6, arg_2)
  return arg_6, arg_7, arg_8


def Func(arg_0: arg_1, arg_2: arg_1) -> arg_1:
  """Maybe broadcasts `from_structure` to `to_structure`.

  If `from_structure` is a singleton, it is tiled to match the structure of
  `to_structure`. Note that the elements in `from_structure` are not copied if
  this tiling occurs.

  Args:
    from_structure: A structure.
    to_structure: A structure.

  Returns:
    new_from_structure: Same structure as `to_structure`.
  """
  arg_3 = tf.nest.flatten(arg_0)
  arg_4 = tf.nest.flatten(arg_2)
  if len(arg_3) == 1:
    arg_3 *= len(arg_4)
  return tf.nest.pack_sequence_as(arg_2, arg_3)


def Func(arg_0: arg_1,
                          arg_2: arg_3,
                          arg_4: arg_5 = None
                         ) -> Union[arg_1, Tuple[arg_1, arg_5]]:
  """Transforms a log-prob function using a bijector.

  This takes a log-prob function and creates a new log-prob function that now
  takes takes state in the domain of the bijector, forward transforms that state
  and calls the original log-prob function. It then returns the log-probability
  that correctly accounts for this transformation.

  The forward-transformed state is pre-pended to the original log-prob
  function's extra returns and returned as the new extra return.

  For convenience you can also pass the initial state (in the original space),
  and this function will return the inverse transformed as the 2nd return value.
  You'd use this to initialize MCMC operators that operate in the transformed
  space.

  Args:
    log_prob_fn: Log prob fn.
    bijector: Bijector(s), must be of the same structure as the `log_prob_fn`
      inputs.
    init_state: Initial state, in the original space.

  Returns:
    transformed_log_prob_fn: Transformed log prob fn.
    transformed_init_state: If `init_state` is provided. Initial state in the
      transformed space.
  """

  def wrapper(*arg_6):
    """Transformed wrapper."""
    arg_7 = arg_2

    arg_6 = tf.nest.map_structure(lambda x: 0. + x, arg_6)
    if len(arg_6) == 1:
      arg_6 = arg_6[0]
    elif isinstance(arg_7, list):
      arg_7 = tuple(arg_7)

    arg_8 = tf.nest.map_structure(lambda b, x: b.forward(x),
                                                arg_7, arg_6)
    arg_8 = arg_8  # type: Tuple[Any]
    arg_9, arg_10 = call_fn(arg_0, arg_8)
    arg_11 = tf.nest.map_structure(
        lambda x: tf.rank(x) - tf.rank(arg_9), arg_6)

    return arg_9 + sum(
        tf.nest.flatten(
            tf.nest.map_structure(
                lambda b, x, e: b.forward_log_det_jacobian(x, arg_11=e),
                arg_7, arg_6, arg_11))), [arg_8, arg_10]

  if arg_4 is None:
    return wrapper
  else:
    return wrapper, tf.nest.map_structure(lambda b, s: b.inverse(s), arg_2,
                                          arg_4)


def Func(arg_0: arg_1,
                  arg_2: arg_3, arg_4: arg_5,
                  arg_6: arg_5
                 ) -> Tuple[arg_1, LeapFrogStepExtras]:
  """Leapfrog `TransitionOperator`.

  Args:
    Func_state: LeapFrogStepState.
    step_size: Step size, structure broadcastable to the `target_log_prob_fn`
      state.
    target_log_prob_fn: Target log prob fn.
    kinetic_energy_fn: Kinetic energy fn.

  Returns:
    Func_state: LeapFrogStepState.
    Func_extras: LeapFrogStepExtras.
  """
  arg_7 = arg_0.state
  arg_8 = arg_0.state_grads
  arg_9 = arg_0.momentum
  arg_2 = maybe_broadcast_structure(arg_2, arg_7)

  arg_7 = tf.nest.map_structure(tf.convert_to_tensor, arg_7)
  arg_9 = tf.nest.map_structure(tf.convert_to_tensor, arg_9)
  arg_7 = tf.nest.map_structure(tf.convert_to_tensor, arg_7)

  if arg_8 is None:
    arg_10, arg_10, arg_8 = call_and_grads(arg_4, arg_7)
  else:
    arg_8 = tf.nest.map_structure(tf.convert_to_tensor, arg_8)

  arg_9 = tf.nest.map_structure(lambda m, sg, s: m + 0.5 * sg * s, arg_9,
                                   arg_8, arg_2)

  arg_11, arg_12, arg_13 = call_and_grads(
      arg_6, arg_9)

  arg_7 = tf.nest.map_structure(lambda x, mg, s: x + mg * s, arg_7,
                                arg_13, arg_2)

  arg_14, arg_15, arg_8 = call_and_grads(
      arg_4, arg_7)

  arg_9 = tf.nest.map_structure(lambda m, sg, s: m + 0.5 * sg * s, arg_9,
                                   arg_8, arg_2)

  return arg_1(arg_7, arg_8, arg_9), LeapFrogStepExtras(
      arg_14, arg_15, arg_11, arg_12)


def Func(arg_0: arg_1,
                             arg_2: arg_1,
                             arg_3: arg_4,
                             arg_5=None) -> Tuple[arg_1, tf.Tensor, tf.Tensor]:
  """Metropolis-Hastings step.

  This probabilistically chooses between `current_state` and `proposed_state`
  based on the `energy_change` so as to preserve detailed balance.

  Energy change is the negative of `log_accept_ratio`.

  Args:
    current_state: Current state.
    proposed_state: Proposed state.
    energy_change: E(proposed_state) - E(previous_state).
    seed: For reproducibility.

  Returns:
    new_state: The chosen state.
    is_accepted: Whether the proposed state was accepted.
    log_uniform: The random number that was used to select between the two
      states.
  """
  arg_6 = tf.nest.flatten(arg_0)
  arg_7 = nest.flatten_up_to(arg_0, arg_2)
  # Impute the None's in the current state.
  arg_6 = [
      p if c is None else c for p, c in zip(arg_7, arg_6)
  ]
  arg_0 = tf.nest.pack_sequence_as(arg_0, arg_6)

  arg_0 = tf.nest.map_structure(tf.convert_to_tensor, arg_0)
  arg_2 = tf.nest.map_structure(tf.convert_to_tensor, arg_2)
  arg_3 = tf.convert_to_tensor(value=arg_3)

  arg_8 = -arg_3

  arg_9 = tf.math.log(
      tf.random.uniform(
          shape=tf.shape(input=arg_8),
          dtype=arg_8.dtype.base_dtype,
          arg_5=arg_5))
  arg_10 = arg_9 < arg_8

  arg_11 = mcmc_util.choose(
      arg_10, arg_2, arg_0, name='choose_next_state')
  return arg_11, arg_10, arg_9


def Func(arg_0, arg_1, arg_2, arg_3,
                             arg_4, arg_5, arg_6, arg_7,
                             arg_8):
    """Construct `scale` from various components.

    Args:
      identity_multiplier: floating point rank 0 `Tensor` representing a scaling
        done to the identity matrix.
      diag: Floating-point `Tensor` representing the diagonal matrix.`diag` has
        shape `[N1, N2, ...  k]`, which represents a k x k diagonal matrix.
      tril: Floating-point `Tensor` representing the lower triangular matrix.
       `tril` has shape `[N1, N2, ...  k, k]`, which represents a k x k lower
       triangular matrix.
      perturb_diag: Floating-point `Tensor` representing the diagonal matrix of
        the low rank update.
      perturb_factor: Floating-point `Tensor` representing factor matrix.
      shift: Floating-point `Tensor` representing `shift in `scale @ X + shift`.
      validate_args: Python `bool` indicating whether arguments should be
        checked for correctness.
      dtype: `DType` for arg `Tensor` conversions.

    Returns:
      scale. In the case of scaling by a constant, scale is a
      floating point `Tensor`. Otherwise, scale is a `LinearOperator`.

    Raises:
      ValueError: if all of `tril`, `diag` and `identity_multiplier` are `None`.
    """
    arg_1 = _as_tensor(arg_1, "identity_multiplier",
                                     arg_8)
    arg_2 = _as_tensor(arg_2, "diag", arg_8)
    arg_3 = _as_tensor(arg_3, "tril", arg_8)
    arg_4 = _as_tensor(arg_4, "perturb_diag", arg_8)
    arg_5 = _as_tensor(arg_5, "perturb_factor", arg_8)

    # If possible, use the low rank update to infer the shape of
    # the identity matrix, when scale represents a scaled identity matrix
    # with a low rank update.
    arg_9 = None
    if arg_5 is not None:
      arg_9 = distribution_util.dimension_size(arg_5, axis=-2)

    if arg_0._is_only_identity_multiplier:
      if arg_7:
        return distribution_util.with_dependencies([
            assert_util.assert_none_equal(
                arg_1, tf.zeros([], arg_1.dtype),
                ["identity_multiplier should be non-zero."])
        ], arg_1)
      return arg_1

    arg_10 = distribution_util.make_tril_scale(
        loc=arg_6,
        scale_tril=arg_3,
        scale_diag=arg_2,
        scale_identity_multiplier=arg_1,
        arg_7=arg_7,
        assert_positive=False,
        arg_9=arg_9)

    if arg_5 is not None:
      return tf.linalg.LinearOperatorLowRankUpdate(
          arg_10,
          u=arg_5,
          diag_update=arg_4,
          is_diag_update_positive=arg_4 is None,
          is_non_singular=True,  # Implied by is_positive_definite=True.
          is_self_adjoint=True,
          is_positive_definite=True,
          is_square=True)

    return arg_10


def Func(arg_0=1., arg_1=None):
  """Returns a callable that adds a random normal perturbation to the input.

  This function returns a callable that accepts a Python `list` of `Tensor`s of
  any shapes and `dtypes`  representing the state parts of the `current_state`
  and a random seed. The supplied argument `scale` must be a `Tensor` or Python
  `list` of `Tensor`s representing the scale of the generated
  proposal. `scale` must broadcast with the state parts of `current_state`.
  The callable adds a sample from a zero-mean normal distribution with the
  supplied scales to each state part and returns a same-type `list` of `Tensor`s
  as the state parts of `current_state`.

  Args:
    scale: a `Tensor` or Python `list` of `Tensor`s of any shapes and `dtypes`
      controlling the scale of the normal proposal distribution.
    name: Python `str` name prefixed to Ops created by this function.
        Default value: 'Func'.

  Returns:
    Func: A callable accepting a Python `list` of `Tensor`s
      representing the state parts of the `current_state` and an `int`
      representing the random seed to be used to generate the proposal. The
      callable returns the same-type `list` of `Tensor`s as the input and
      represents the proposal for the RWM algorithm.
  """
  def _fn(arg_2, arg_3):
    """Adds a normal perturbation to the input state.

    Args:
      state_parts: A list of `Tensor`s of any shape and real dtype representing
        the state parts of the `current_state` of the Markov chain.
      seed: `int` or None. The random seed for this `Op`. If `None`, no seed is
        applied.
        Default value: `None`.

    Returns:
      perturbed_state_parts: A Python `list` of The `Tensor`s. Has the same
        shape and type as the `state_parts`.

    Raises:
      ValueError: if `scale` does not broadcast with `state_parts`.
    """
    with tf.compat.v1.name_scope(
        arg_1, 'Func', values=[arg_2, arg_0, arg_3]):
      arg_4 = arg_0 if mcmc_util.is_list_like(arg_0) else [arg_0]
      if len(arg_4) == 1:
        arg_4 *= len(arg_2)
      if len(arg_2) != len(arg_4):
        raise ValueError('`scale` must broadcast with `state_parts`.')
      arg_5 = distributions.SeedStream(arg_3, salt='RandomWalkNormalFn')
      arg_6 = [
          tf.random.normal(
              mean=state_part,
              stddev=scale_part,
              shape=tf.shape(input=state_part),
              dtype=state_part.dtype.base_dtype,
              arg_3=arg_5())
          for scale_part, state_part in zip(arg_4, arg_2)
      ]

      return arg_6
  return _fn


def Func(arg_0=1., arg_1=None):
  """Returns a callable that adds a random uniform perturbation to the input.

  For more details on `Func`, see
  `random_walk_normal_fn`. `scale` might
  be a `Tensor` or a list of `Tensor`s that should broadcast with state parts
  of the `current_state`. The generated uniform perturbation is sampled as a
  uniform point on the rectangle `[-scale, scale]`.

  Args:
    scale: a `Tensor` or Python `list` of `Tensor`s of any shapes and `dtypes`
      controlling the upper and lower bound of the uniform proposal
      distribution.
    name: Python `str` name prefixed to Ops created by this function.
        Default value: 'Func'.

  Returns:
    Func: A callable accepting a Python `list` of `Tensor`s
      representing the state parts of the `current_state` and an `int`
      representing the random seed used to generate the proposal. The callable
      returns the same-type `list` of `Tensor`s as the input and represents the
      proposal for the RWM algorithm.
  """
  def _fn(arg_2, arg_3):
    """Adds a uniform perturbation to the input state.

    Args:
      state_parts: A list of `Tensor`s of any shape and real dtype representing
        the state parts of the `current_state` of the Markov chain.
      seed: `int` or None. The random seed for this `Op`. If `None`, no seed is
        applied.
        Default value: `None`.

    Returns:
      perturbed_state_parts: A Python `list` of The `Tensor`s. Has the same
        shape and type as the `state_parts`.

    Raises:
      ValueError: if `scale` does not broadcast with `state_parts`.
    """
    with tf.compat.v1.name_scope(
        arg_1, 'Func', values=[arg_2, arg_0, arg_3]):
      arg_4 = arg_0 if mcmc_util.is_list_like(arg_0) else [arg_0]
      if len(arg_4) == 1:
        arg_4 *= len(arg_2)
      if len(arg_2) != len(arg_4):
        raise ValueError('`scale` must broadcast with `state_parts`.')
      arg_5 = distributions.SeedStream(arg_3, salt='RandomWalkUniformFn')
      arg_6 = [
          tf.random.uniform(
              minval=state_part - scale_part,
              maxval=state_part + scale_part,
              shape=tf.shape(input=state_part),
              dtype=state_part.dtype.base_dtype,
              arg_3=arg_5())
          for scale_part, state_part in zip(arg_4, arg_2)
      ]
      return arg_6
  return _fn


def Func(arg_0, arg_1):
    """Get a list of num_components batchwise probabilities."""
    arg_2 = tf.nn.log_softmax if arg_1 else tf.nn.softmax
    arg_3 = arg_2(arg_0.cat.logits)
    arg_3 = tf.unstack(arg_3, num=arg_0.num_components, axis=-1)
    return arg_3


def Func(arg_0, arg_1, arg_2, arg_3):
  """Validate `outcomes`, `logits` and `probs`'s shapes."""
  arg_4 = []

  def validate_equal_last_dim(arg_5, arg_6, arg_7):
    if arg_5.shape.is_fully_defined() and arg_6.shape.is_fully_defined():
      if arg_5.shape[-1] != arg_6.shape[-1]:
        raise ValueError(arg_7)
    elif arg_3:
      arg_4.append(
          tf.compat.v1.assert_equal(
              tf.shape(input=arg_5)[-1],
              tf.shape(input=arg_6)[-1],
              arg_7=arg_7))

  if arg_1 is not None:
    validate_equal_last_dim(
        arg_0,
        arg_1,
        arg_7='Last dimension of outcomes and logits must be equal size.')
  if arg_2 is not None:
    validate_equal_last_dim(
        arg_0,
        arg_2,
        arg_7='Last dimension of outcomes and probs must be equal size.')

  arg_7 = 'Rank of outcomes must be 1.'
  if arg_0.shape.ndims is not None:
    if arg_0.shape.ndims != 1:
      raise ValueError(arg_7)
  elif arg_3:
    arg_4.append(tf.compat.v1.assert_rank(arg_0, 1, arg_7=arg_7))

  arg_7 = 'Size of outcomes must be greater than 0.'
  if arg_0.shape.num_elements() is not None:
    if arg_0.shape.num_elements() == 0:
      raise ValueError(arg_7)
  elif arg_3:
    arg_4.append(
        tf.compat.v1.assert_greater(
            tf.size(input=arg_0), 0, arg_7=arg_7))

  if arg_3:
    arg_4.append(
        tf.compat.v1.assert_equal(
            tf.math.is_strictly_increasing(arg_0),
            True,
            arg_7='outcomes is not strictly increasing.'))

  return arg_4


def Func(arg_0):
  """Bayesian logistic regression, which returns labels given features."""
  arg_1 = ed.MultivariateNormalDiag(
      loc=tf.zeros(arg_0.shape[1]), name="coeffs")
  arg_2 = ed.Bernoulli(
      logits=tf.tensordot(arg_0, arg_1, [[1], [0]]), name="labels")
  return arg_2


def Func():
  """Builds the Covertype data set."""
  import sklearn.datasets  # pylint: disable=g-import-not-at-top
  arg_0 = sklearn.datasets.covtype.fetch_covtype()
  arg_1 = arg_0.data
  arg_2 = arg_0.target

  # Normalize features and append a column of ones for the intercept.
  arg_1 -= arg_1.mean(0)
  arg_1 /= arg_1.std(0)
  arg_1 = np.hstack([arg_1, np.ones([arg_1.shape[0], 1])])
  arg_1 = tf.cast(arg_1, dtype=tf.float32)

  # Binarize outcomes on whether it is a specific category.
  arg_3, arg_4 = np.unique(arg_2, return_counts=True)
  arg_5 = np.argmax(arg_4)
  arg_2 = (arg_2 == arg_5)
  arg_2 = tf.cast(arg_2, dtype=tf.int32)
  return arg_1, arg_2


def Func(arg_0, arg_1=0, arg_2=False, arg_3=None):
  """Cholesky factor of the covariance matrix of vector-variate random samples.

  This function can be use to fit a multivariate normal to data.

  ```python
  tf.enable_eager_execution()
  import tensorflow_probability as tfp
  tfd = tfp.distributions

  # Assume data.shape = (1000, 2).  1000 samples of a random variable in R^2.
  observed_data = read_data_samples(...)

  # The mean is easy
  mu = tf.reduce_mean(observed_data, axis=0)

  # Get the scale matrix
  L = tfp.stats.Func(observed_data)

  # Make the best fit multivariate normal (under maximum likelihood condition).
  mvn = tfd.MultivariateNormalTriL(loc=mu, scale_tril=L)

  # Plot contours of the pdf.
  xs, ys = tf.meshgrid(
      tf.linspace(-5., 5., 50), tf.linspace(-5., 5., 50), indexing='ij')
  xy = tf.stack((tf.reshape(xs, [-1]), tf.reshape(ys, [-1])), axis=-1)
  pdf = tf.reshape(mvn.prob(xy), (50, 50))
  CS = plt.contour(xs, ys, pdf, 10)
  plt.clabel(CS, inline=1, fontsize=10)
  ```

  Why does this work?
  Given vector-variate random variables `X = (X1, ..., Xd)`, one may obtain the
  sample covariance matrix in `R^{d x d}` (see `tfp.stats.covariance`).

  The [Cholesky factor](https://en.wikipedia.org/wiki/Cholesky_decomposition)
  of this matrix is analogous to standard deviation for scalar random variables:
  Suppose `X` has covariance matrix `C`, with Cholesky factorization `C = L L^T`
  Then multiplying a vector of iid random variables which have unit variance by
  `L` produces a vector with covariance `L L^T`, which is the same as `X`.

  ```python
  observed_data = read_data_samples(...)
  L = tfp.stats.Func(observed_data, sample_axis=0)

  # Make fake_data with the same covariance as observed_data.
  uncorrelated_normal = tf.random_normal(shape=(500, 10))
  fake_data = tf.linalg.matvec(L, uncorrelated_normal)
  ```

  Args:
    x:  Numeric `Tensor`.  The rightmost dimension of `x` indexes events. E.g.
      dimensions of a random vector.
    sample_axis: Scalar or vector `Tensor` designating axis holding samples.
      Default value: `0` (leftmost dimension). Cannot be the rightmost dimension
        (since this indexes events).
    keepdims:  Boolean.  Whether to keep the sample axis as singletons.
    name: Python `str` name prefixed to Ops created by this function.
          Default value: `None` (i.e., `'covariance'`).

  Returns:
    chol:  `Tensor` of same `dtype` as `x`.  The last two dimensions hold
      lower triangular matrices (the Cholesky factors).
  """
  with tf.compat.v1.name_scope(
      arg_3, 'Func', values=[arg_0, arg_1]):
    arg_1 = tf.convert_to_tensor(value=arg_1, dtype=tf.int32)
    arg_4 = covariance(
        arg_0, arg_1=arg_1, event_axis=-1, arg_2=arg_2)
    return tf.linalg.cholesky(arg_4)


def Func(arg_0, arg_1=0, arg_2=False, arg_3=None):
  """Estimate standard deviation using samples.

  Given `N` samples of scalar valued random variable `X`, standard deviation may
  be estimated as

  ```none
  Stddev[X] := Sqrt[Var[X]],
  Var[X] := N^{-1} sum_{n=1}^N (X_n - Xbar) Conj{(X_n - Xbar)},
  Xbar := N^{-1} sum_{n=1}^N X_n
  ```

  ```python
  x = tf.random_normal(shape=(100, 2, 3))

  # Func[i, j] is the sample standard deviation of the (i, j) batch member.
  Func = tfp.stats.Func(x, sample_axis=0)
  ```

  Scaling a unit normal by a standard deviation produces normal samples
  with that standard deviation.

  ```python
  observed_data = read_data_samples(...)
  Func = tfp.stats.Func(observed_data)

  # Make fake_data with the same standard deviation as observed_data.
  fake_data = Func * tf.random_normal(shape=(100,))
  ```

  Notice we divide by `N` (the numpy default), which does not create `NaN`
  when `N = 1`, but is slightly biased.

  Args:
    x:  A numeric `Tensor` holding samples.
    sample_axis: Scalar or vector `Tensor` designating axis holding samples, or
      `None` (meaning all axis hold samples).
      Default value: `0` (leftmost dimension).
    keepdims:  Boolean.  Whether to keep the sample axis as singletons.
    name: Python `str` name prefixed to Ops created by this function.
          Default value: `None` (i.e., `'Func'`).

  Returns:
    Func: A `Tensor` of same `dtype` as the `x`, and rank equal to
      `rank(x) - len(sample_axis)`
  """
  with tf.compat.v1.name_scope(arg_3, 'Func', values=[arg_0, arg_1]):
    return tf.sqrt(variance(arg_0, arg_1=arg_1, arg_2=arg_2))


def Func(arg_0, arg_1):
  """Rectify possibly negatively axis. Prefer return Python list."""
  arg_0 = _make_list_or_1d_tensor(arg_0)

  arg_1 = tf.convert_to_tensor(value=arg_1, name='ndims', dtype=tf.int32)
  arg_2 = tf.get_static_value(arg_1)

  if _is_list_like(arg_0) and arg_2 is not None:
    # Static case
    arg_3 = []
    for arg_4 in arg_0:
      if arg_4 < 0:
        arg_4 = arg_2 + arg_4
      arg_3.append(arg_4)
  else:
    # Dynamic case
    arg_0 = tf.convert_to_tensor(value=arg_0, name='axis', dtype=tf.int32)
    arg_3 = tf.where(arg_0 >= 0, arg_0, arg_0 + arg_1)

  return arg_3


def Func(arg_0, arg_1):
  """A version of squeeze that works with dynamic axis."""
  arg_0 = tf.convert_to_tensor(value=arg_0, name='x')
  if arg_1 is None:
    return tf.squeeze(arg_0, arg_1=None)
  arg_1 = tf.convert_to_tensor(value=arg_1, name='axis', dtype=tf.int32)
  arg_1 += tf.zeros([1], dtype=arg_1.dtype)  # Make axis at least 1d.
  arg_2, arg_3 = tf.compat.v1.setdiff1d(tf.range(0, tf.rank(arg_0)), arg_1)
  return tf.reshape(arg_0, tf.gather(tf.shape(input=arg_0), arg_2))


def Func(arg_0, arg_1):
    """Standardize input `x` to a unit normal."""
    with tf.name_scope("standardize"):
      return (arg_1 - arg_0.loc) / arg_0.scale


def Func(arg_0,
                           arg_1=None,
                           arg_2=None,
                           arg_3=arg_4.float32,
                           arg_6=True,
                           arg_7=None,
                           arg_8=None):
  r"""Returns a sample from the `dim` dimensional Halton sequence.

  Warning: The sequence elements take values only between 0 and 1. Care must be
  taken to appropriately transform the domain of a function if it differs from
  the unit cube before evaluating integrals using Halton samples. It is also
  important to remember that quasi-random numbers without randomization are not
  a replacement for pseudo-random numbers in every context. Quasi random numbers
  are completely deterministic and typically have significant negative
  autocorrelation unless randomization is used.

  Computes the members of the low discrepancy Halton sequence in dimension
  `dim`. The `dim`-dimensional sequence takes values in the unit hypercube in
  `dim` dimensions. Currently, only dimensions up to 1000 are supported. The
  prime base for the k-th axes is the k-th prime starting from 2. For example,
  if `dim` = 3, then the bases will be [2, 3, 5] respectively and the first
  element of the non-randomized sequence will be: [0.5, 0.333, 0.2]. For a more
  complete description of the Halton sequences see
  [here](https://en.wikipedia.org/wiki/Halton_sequence). For low discrepancy
  sequences and their applications see
  [here](https://en.wikipedia.org/wiki/Low-discrepancy_sequence).

  If `randomized` is true, this function produces a scrambled version of the
  Halton sequence introduced by [Owen (2017)][1]. For the advantages of
  randomization of low discrepancy sequences see [here](
  https://en.wikipedia.org/wiki/Quasi-Monte_Carlo_method#Randomization_of_quasi-Monte_Carlo).

  The number of samples produced is controlled by the `num_results` and
  `sequence_indices` parameters. The user must supply either `num_results` or
  `sequence_indices` but not both.
  The former is the number of samples to produce starting from the first
  element. If `sequence_indices` is given instead, the specified elements of
  the sequence are generated. For example, sequence_indices=tf.range(10) is
  equivalent to specifying n=10.

  #### Examples

  ```python
  import tensorflow as tf
  import tensorflow_probability as tfp

  # Produce the first 1000 members of the Halton sequence in 3 dimensions.
  num_results = 1000
  dim = 3
  sample = tfp.mcmc.Func(
    dim,
    num_results=num_results,
    seed=127)

  # Evaluate the integral of x_1 * x_2^2 * x_3^3  over the three dimensional
  # hypercube.
  powers = tf.range(1.0, limit=dim + 1)
  integral = tf.reduce_mean(tf.reduce_prod(sample ** powers, axis=-1))
  true_value = 1.0 / tf.reduce_prod(powers + 1.0)
  with tf.Session() as session:
    values = session.run((integral, true_value))

  # Produces a relative absolute error of 1.7%.
  print ("Estimated: %f, True Value: %f" % values)

  # Now skip the first 1000 samples and recompute the integral with the next
  # thousand samples. The sequence_indices argument can be used to do this.

  sequence_indices = tf.range(start=1000, limit=1000 + num_results,
                              dtype=tf.int32)
  sample_leaped = tfp.mcmc.Func(
      dim,
      sequence_indices=sequence_indices,
      seed=111217)

  integral_leaped = tf.reduce_mean(tf.reduce_prod(sample_leaped ** powers,
                                                  axis=-1))
  with tf.Session() as session:
    values = session.run((integral_leaped, true_value))
  # Now produces a relative absolute error of 0.05%.
  print ("Leaped Estimated: %f, True Value: %f" % values)
  ```

  Args:
    dim: Positive Python `int` representing each sample's `event_size.` Must
      not be greater than 1000.
    num_results: (Optional) Positive scalar `Tensor` of dtype int32. The number
      of samples to generate. Either this parameter or sequence_indices must
      be specified but not both. If this parameter is None, then the behaviour
      is determined by the `sequence_indices`.
      Default value: `None`.
    sequence_indices: (Optional) `Tensor` of dtype int32 and rank 1. The
      elements of the sequence to compute specified by their position in the
      sequence. The entries index into the Halton sequence starting with 0 and
      hence, must be whole numbers. For example, sequence_indices=[0, 5, 6] will
      produce the first, sixth and seventh elements of the sequence. If this
      parameter is None, then the `num_results` parameter must be specified
      which gives the number of desired samples starting from the first sample.
      Default value: `None`.
    dtype: (Optional) The dtype of the sample. One of: `float16`, `float32` or
      `float64`.
      Default value: `tf.float32`.
    randomized: (Optional) bool indicating whether to produce a randomized
      Halton sequence. If True, applies the randomization described in
      [Owen (2017)][1].
      Default value: `True`.
    seed: (Optional) Python integer to seed the random number generator. Only
      used if `randomized` is True. If not supplied and `randomized` is True,
      no seed is set.
      Default value: `None`.
    name:  (Optional) Python `str` describing ops managed by this function. If
      not supplied the name of this function is used.
      Default value: "Func".

  Returns:
    halton_elements: Elements of the Halton sequence. `Tensor` of supplied dtype
      and `shape` `[num_results, dim]` if `num_results` was specified or shape
      `[s, dim]` where s is the size of `sequence_indices` if `sequence_indices`
      were specified.

  Raises:
    ValueError: if both `sequence_indices` and `num_results` were specified or
      if dimension `dim` is less than 1 or greater than 1000.

  #### References

  [1]: Art B. Owen. A randomized Halton algorithm in R. _arXiv preprint
       arXiv:1706.02808_, 2017. https://arxiv.org/abs/1706.02808
  """
  if arg_0 < 1 or arg_0 > _MAX_DIMENSION:
    raise ValueError(
        'Dimension must be between 1 and {}. Supplied {}'.format(_MAX_DIMENSION,
                                                                 arg_0))
  if (arg_1 is None) == (arg_2 is None):
    raise ValueError('Either `num_results` or `sequence_indices` must be'
                     ' specified but not both.')

  if not arg_3.is_floating:
    raise ValueError('dtype must be of `float`-type')

  with arg_4.compat.v1.name_scope(
      arg_8, 'sample', values=[arg_1, arg_2]):
    # Here and in the following, the shape layout is as follows:
    # [sample dimension, event dimension, coefficient dimension].
    # The coefficient dimension is an intermediate axes which will hold the
    # weights of the starting integer when expressed in the (prime) base for
    # an event dimension.
    if arg_1 is not None:
      arg_1 = arg_4.convert_to_tensor(value=arg_1)
    if arg_2 is not None:
      arg_2 = arg_4.convert_to_tensor(value=arg_2)
    arg_9 = _get_indices(arg_1, arg_2, arg_3)
    arg_10 = arg_4.constant(_PRIMES[0:arg_0], arg_3=arg_3, shape=[arg_0, 1])

    arg_11 = _base_expansion_size(
        arg_4.reduce_max(input_tensor=arg_9), arg_10)

    arg_12 = arg_4.reduce_max(input_tensor=arg_11)

    # The powers of the radixes that we will need. Note that there is a bit
    # of an excess here. Suppose we need the place value coefficients of 7
    # in base 2 and 3. For 2, we will have 3 digits but we only need 2 digits
    # for base 3. However, we can only create rectangular tensors so we
    # store both expansions in a [2, 3] tensor. This leads to the problem that
    # we might end up attempting to raise large numbers to large powers. For
    # example, base 2 expansion of 1024 has 10 digits. If we were in 10
    # dimensions, then the 10th prime (29) we will end up computing 29^10 even
    # though we don't need it. We avoid this by setting the exponents for each
    # axes to 0 beyond the maximum value needed for that dimension.
    arg_13 = arg_4.tile([arg_4.range(arg_12)], [arg_0, 1])

    # The mask is true for those coefficients that are irrelevant.
    arg_14 = arg_13 >= arg_11
    arg_15 = arg_4.where(
        arg_14,
        arg_4.zeros_like(arg_13),
        arg_13)
    arg_16 = arg_10 ** arg_15
    # The following computes the base b expansion of the indices. Suppose,
    # x = a0 + a1*b + a2*b^2 + ... Then, performing a floor div of x with
    # the vector (1, b, b^2, b^3, ...) will produce
    # (a0 + s1 * b, a1 + s2 * b, ...) where s_i are coefficients we don't care
    # about. Noting that all a_i < b by definition of place value expansion,
    # we see that taking the elements mod b of the above vector produces the
    # place value expansion coefficients.
    arg_17 = arg_4.math.floordiv(arg_9, arg_16)
    arg_17 *= 1. - arg_4.cast(arg_14, arg_3)
    arg_17 %= arg_10
    if not arg_6:
      arg_17 /= arg_10
      return arg_4.reduce_sum(input_tensor=arg_17 / arg_16, axis=-1)
    arg_18 = distributions.SeedStream(arg_7, salt='MCMCSampleHaltonSequence')
    arg_17 = _randomize(arg_17, arg_10, arg_7=arg_18())
    # Remove the contribution from randomizing the trailing zero for the
    # axes where max_size_by_axes < max_size. This will be accounted
    # for separately below (using zero_correction).
    arg_17 *= 1. - arg_4.cast(arg_14, arg_3)
    arg_17 /= arg_10
    arg_19 = arg_4.reduce_sum(input_tensor=arg_17 / arg_16, axis=-1)

    # The randomization used in Owen (2017) does not leave 0 invariant. While
    # we have accounted for the randomization of the first `max_size_by_axes`
    # coefficients, we still need to correct for the trailing zeros. Luckily,
    # this is equivalent to adding a uniform random value scaled so the first
    # `max_size_by_axes` coefficients are zero. The following statements perform
    # this correction.
    arg_20 = arg_4.random.uniform([arg_0, 1], arg_7=arg_18(), arg_3=arg_3)
    arg_20 /= arg_10 ** arg_11
    return arg_19 + arg_4.reshape(arg_20, [-1])


def Func(arg_0, arg_1, arg_2=None):
  """Uniform iid sample from the space of permutations.

  Draws a sample of size `num_results` from the group of permutations of degrees
  specified by the `dims` tensor. These are packed together into one tensor
  such that each row is one sample from each of the dimensions in `dims`. For
  example, if dims = [2,3] and num_results = 2, the result is a tensor of shape
  [2, 2 + 3] and the first row of the result might look like:
  [1, 0, 2, 0, 1]. The first two elements are a permutation over 2 elements
  while the next three are a permutation over 3 elements.

  Args:
    num_results: A positive scalar `Tensor` of integral type. The number of
      draws from the discrete uniform distribution over the permutation groups.
    dims: A 1D `Tensor` of the same dtype as `num_results`. The degree of the
      permutation groups from which to sample.
    seed: (Optional) Python integer to seed the random number generator.

  Returns:
    permutations: A `Tensor` of shape `[num_results, sum(dims)]` and the same
    dtype as `dims`.
  """
  arg_3 = tf.range(arg_0)
  arg_4 = distributions.SeedStream(arg_2, salt='MCMCSampleHaltonSequence3')
  def generate_one(arg_5):
    arg_2 = arg_4()
    arg_6 = lambda _: tf.random.shuffle(tf.range(arg_5), arg_2=arg_2)
    return tf.map_fn(
        arg_6,
        arg_3,
        parallel_iterations=1 if arg_2 is not None else 10)
  return tf.concat([generate_one(arg_5) for arg_5 in tf.unstack(arg_1)],
                   axis=-1)


def Func(arg_0, arg_1, arg_2, arg_3=None):
  """Generates starting points for the Halton sequence procedure.

  The k'th element of the sequence is generated starting from a positive integer
  which must be distinct for each `k`. It is conventional to choose the starting
  point as `k` itself (or `k+1` if k is zero based). This function generates
  the starting integers for the required elements and reshapes the result for
  later use.

  Args:
    num_results: Positive scalar `Tensor` of dtype int32. The number of samples
      to generate. If this parameter is supplied, then `sequence_indices`
      should be None.
    sequence_indices: `Tensor` of dtype int32 and rank 1. The entries
      index into the Halton sequence starting with 0 and hence, must be whole
      numbers. For example, sequence_indices=[0, 5, 6] will produce the first,
      sixth and seventh elements of the sequence. If this parameter is not None
      then `n` must be None.
    dtype: The dtype of the sample. One of `float32` or `float64`.
      Default is `float32`.
    name: Python `str` name which describes ops created by this function.

  Returns:
    indices: `Tensor` of dtype `dtype` and shape = `[n, 1, 1]`.
  """
  with tf.compat.v1.name_scope(arg_3, 'Func',
                               [arg_0, arg_1]):
    if arg_1 is None:
      arg_0 = tf.cast(arg_0, arg_2=arg_2)
      arg_1 = tf.range(arg_0, arg_2=arg_2)
    else:
      arg_1 = tf.cast(arg_1, arg_2)

    # Shift the indices so they are 1 based.
    arg_4 = arg_1 + 1

    # Reshape to make space for the event dimension and the place value
    # coefficients.
    return tf.reshape(arg_4, [-1, 1, 1])


def Func(arg_0, arg_1):
  """Computes the number of terms in the place value expansion.

  Let num = a0 + a1 b + a2 b^2 + ... ak b^k be the place value expansion of
  `num` in base b (ak <> 0). This function computes and returns `k+1` for each
  base `b` specified in `bases`.

  This can be inferred from the base `b` logarithm of `num` as follows:
    $$k = Floor(log_b (num)) + 1  = Floor( log(num) / log(b)) + 1$$

  Args:
    num: Scalar `Tensor` of dtype either `float32` or `float64`. The number to
      compute the base expansion size of.
    bases: `Tensor` of the same dtype as num. The bases to compute the size
      against.

  Returns:
    Tensor of same dtype and shape as `bases` containing the size of num when
    written in that base.
  """
  return tf.floor(tf.math.log(arg_0) / tf.math.log(arg_1)) + 1


def Func(arg_0):
  # Based on
  # https://stackoverflow.com/questions/2068372/fastest-way-to-list-all-primes-below-n-in-python/3035188#3035188
  """Returns sorted array of primes such that `2 <= prime < n`."""
  arg_1 = np.array((2, 3, 5))
  if arg_0 <= 6:
    return arg_1[arg_1 < arg_0]
  arg_2 = np.ones(arg_0 // 3 + (arg_0 % 6 == 2), dtype=np.bool)
  arg_2[0] = False
  arg_3 = int(arg_0 ** 0.5) // 3 + 1
  for arg_4 in range(arg_3):
    if not arg_2[arg_4]:
      continue
    arg_5 = 3 * arg_4 + 1 | 1
    arg_2[arg_5 ** 2 // 3::2 * arg_5] = False
    arg_2[(arg_5 ** 2 + 4 * arg_5 - 2 * arg_5 * (arg_4 & 1)) // 3::2 * arg_5] = False
  return np.r_[2, 3, 3 * np.nonzero(arg_2)[0] + 1 | 1]


def Func(arg_0):
  """Returns the machine epsilon for the supplied dtype."""
  if isinstance(arg_0, tf.DType):
    arg_0 = arg_0.as_numpy_dtype()
  return np.finfo(arg_0).eps


def Func(arg_0,
                arg_1=None,
                arg_2=None,
                arg_3=None,
                arg_4=None,
                arg_5=1e-6,
                arg_6=0.66,
                arg_7=5.0,
                arg_8=0.1,
                arg_9=0.9,
                arg_10=0.1,
                arg_11=50,
                arg_12=None):
  """The Hager Zhang line search algorithm.

  Performs an inexact line search based on the algorithm of
  [Hager and Zhang (2006)][2].
  The univariate objective function `value_and_gradients_function` is typically
  generated by projecting a multivariate objective function along a search
  direction. Suppose the multivariate function to be minimized is
  `g(x1,x2, .. xn)`. Let (d1, d2, ..., dn) be the direction along which we wish
  to perform a line search. Then the projected univariate function to be used
  for line search is

  ```None
    f(a) = g(x1 + d1 * a, x2 + d2 * a, ..., xn + dn * a)
  ```

  The directional derivative along (d1, d2, ..., dn) is needed for this
  procedure. This also corresponds to the derivative of the projected function
  `f(a)` with respect to `a`. Note that this derivative must be negative for
  `a = 0` if the direction is a descent direction.

  The usual stopping criteria for the line search is the satisfaction of the
  (weak) Wolfe conditions. For details of the Wolfe conditions, see
  ref. [3]. On a finite precision machine, the exact Wolfe conditions can
  be difficult to satisfy when one is very close to the minimum and as argued
  by [Hager and Zhang (2005)][1], one can only expect the minimum to be
  determined within square root of machine precision. To improve the situation,
  they propose to replace the Wolfe conditions with an approximate version
  depending on the derivative of the function which is applied only when one
  is very close to the minimum. The following algorithm implements this
  enhanced scheme.

  ### Usage:

  Primary use of line search methods is as an internal component of a class of
  optimization algorithms (called line search based methods as opposed to
  trust region methods). Hence, the end user will typically not want to access
  line search directly. In particular, inexact line search should not be
  confused with a univariate minimization method. The stopping criteria of line
  search is the satisfaction of Wolfe conditions and not the discovery of the
  minimum of the function.

  With this caveat in mind, the following example illustrates the standalone
  usage of the line search.

  ```python
    # Define value and gradient namedtuple
    ValueAndGradient = namedtuple('ValueAndGradient', ['x', 'f', 'df'])
    # Define a quadratic target with minimum at 1.3.
    def value_and_gradients_function(x):
      return ValueAndGradient(x=x, f=(x - 1.3) ** 2, df=2 * (x-1.3))
    # Set initial step size.
    step_size = tf.constant(0.1)
    ls_result = tfp.optimizer.linesearch.Func(
        value_and_gradients_function, initial_step_size=step_size)
    # Evaluate the results.
    with tf.Session() as session:
      results = session.run(ls_result)
      # Ensure convergence.
      assert results.converged
      # If the line search converged, the left and the right ends of the
      # bracketing interval are identical.
      assert results.left.x == result.right.x
      # Print the number of evaluations and the final step size.
      print ("Final Step Size: %f, Evaluations: %d" % (results.left.x,
                                                       results.func_evals))
  ```

  ### References:
  [1]: William Hager, Hongchao Zhang. A new conjugate gradient method with
    guaranteed descent and an efficient line search. SIAM J. Optim., Vol 16. 1,
    pp. 170-172. 2005.
    https://www.math.lsu.edu/~hozhang/papers/cg_descent.pdf

  [2]: William Hager, Hongchao Zhang. Algorithm 851: CG_DESCENT, a conjugate
    gradient method with guaranteed descent. ACM Transactions on Mathematical
    Software, Vol 32., 1, pp. 113-137. 2006.
    http://users.clas.ufl.edu/hager/papers/CG/cg_compare.pdf

  [3]: Jorge Nocedal, Stephen Wright. Numerical Optimization. Springer Series in
    Operations Research. pp 33-36. 2006

  Args:
    value_and_gradients_function: A Python callable that accepts a real scalar
      tensor and returns a namedtuple with the fields 'x', 'f', and 'df' that
      correspond to scalar tensors of real dtype containing the point at which
      the function was evaluated, the value of the function, and its
      derivative at that point. The other namedtuple fields, if present,
      should be tensors or sequences (possibly nested) of tensors.
      In usual optimization application, this function would be generated by
      projecting the multivariate objective function along some specific
      direction. The direction is determined by some other procedure but should
      be a descent direction (i.e. the derivative of the projected univariate
      function must be negative at 0.).
      Alternatively, the function may represent the batching of `n` such line
      functions (e.g. projecting a single multivariate objective function along
      `n` distinct directions at once) accepting n points as input, i.e. a
      tensor of shape [n], and the fields 'x', 'f' and 'df' in the returned
      namedtuple should each be a tensor of shape [n], with the corresponding
      input points, function values, and derivatives at those input points.
    initial_step_size: (Optional) Scalar positive `Tensor` of real dtype, or
      a tensor of shape [n] in batching mode. The initial value (or values) to
      try to bracket the minimum. Default is `1.` as a float32.
      Note that this point need not necessarily bracket the minimum for the line
      search to work correctly but the supplied value must be greater than 0.
      A good initial value will make the search converge faster.
    value_at_initial_step: (Optional) The full return value of evaluating
      value_and_gradients_function at initial_step_size, i.e. a namedtuple with
      'x', 'f', 'df', if already known by the caller. If supplied the value of
      `initial_step_size` will be ignored, otherwise the tuple will be computed
      by evaluating value_and_gradients_function.
    value_at_zero: (Optional) The full return value of
      value_and_gradients_function at `0.`, i.e. a namedtuple with
      'x', 'f', 'df', if already known by the caller. If not supplied the tuple
      will be computed by evaluating value_and_gradients_function.
    converged: (Optional) In batching mode a tensor of shape [n], indicating
      batch members which have already converged and no further search should
      be performed. These batch members are also reported as converged in the
      output, and both their `left` and `right` are set to the
      `value_at_initial_step`.
    threshold_use_approximate_wolfe_condition: Scalar positive `Tensor`
      of real dtype. Corresponds to the parameter 'epsilon' in
      [Hager and Zhang (2006)][2]. Used to estimate the
      threshold at which the line search switches to approximate Wolfe
      conditions.
    shrinkage_param: Scalar positive Tensor of real dtype. Must be less than
      `1.`. Corresponds to the parameter `gamma` in
      [Hager and Zhang (2006)][2].
      If the secant**2 step does not shrink the bracketing interval by this
      proportion, a bisection step is performed to reduce the interval width.
    expansion_param: Scalar positive `Tensor` of real dtype. Must be greater
      than `1.`. Used to expand the initial interval in case it does not bracket
      a minimum. Corresponds to `rho` in [Hager and Zhang (2006)][2].
    sufficient_decrease_param: Positive scalar `Tensor` of real dtype.
      Bounded above by the curvature param. Corresponds to `delta` in the
      terminology of [Hager and Zhang (2006)][2].
    curvature_param: Positive scalar `Tensor` of real dtype. Bounded above
      by `1.`. Corresponds to 'sigma' in the terminology of
      [Hager and Zhang (2006)][2].
    step_size_shrink_param: Positive scalar `Tensor` of real dtype. Bounded
      above by `1`. If the supplied step size is too big (i.e. either the
      objective value or the gradient at that point is infinite), this factor
      is used to shrink the step size until it is finite.
    max_iterations: Positive scalar `Tensor` of integral dtype or None. The
      maximum number of iterations to perform in the line search. The number of
      iterations used to bracket the minimum are also counted against this
      parameter.
    name: (Optional) Python str. The name prefixed to the ops created by this
      function. If not supplied, the default name 'Func' is used.

  Returns:
    results: A namedtuple containing the following attributes.
      converged: Boolean `Tensor` of shape [n]. Whether a point satisfying
        Wolfe/Approx wolfe was found.
      failed: Boolean `Tensor` of shape [n]. Whether line search failed e.g.
        if either the objective function or the gradient are not finite at
        an evaluation point.
      iterations: Scalar int32 `Tensor`. Number of line search iterations made.
      func_evals: Scalar int32 `Tensor`. Number of function evaluations made.
      left: A namedtuple, as returned by value_and_gradients_function,
        of the left end point of the final bracketing interval. Values are
        equal to those of `right` on batch members where converged is True.
        Otherwise, it corresponds to the last interval computed.
      right: A namedtuple, as returned by value_and_gradients_function,
        of the right end point of the final bracketing interval. Values are
        equal to those of `left` on batch members where converged is True.
        Otherwise, it corresponds to the last interval computed.
  """
  with tf.compat.v1.name_scope(arg_12, 'Func', [
      arg_1, arg_2, arg_3, arg_4,
      arg_5, arg_6,
      arg_7, arg_8, arg_9]):
    arg_13, arg_14, arg_15, arg_16 = _prepare_args(
        arg_0,
        arg_1,
        arg_2,
        arg_3,
        arg_5)

    arg_17 = (hzl.is_finite(arg_13) & (arg_13.df < 0) &
                    tf.math.is_finite(arg_14.x) & (arg_14.x > 0))

    if arg_4 is None:
      arg_18 = tf.zeros_like(arg_17)  # i.e. all false.
    else:
      arg_18 = tf.convert_to_tensor(value=arg_4)

    arg_19 = ~arg_18 & ~arg_17
    arg_20 = ~arg_18 & arg_17

    # Note: _fix_step_size returns immediately if either all inputs are invalid
    # or none of the active ones need fixing.
    arg_21, arg_22, arg_23 = _fix_step_size(
        arg_0, arg_14, arg_20,
        arg_10)

    arg_24 = HagerZhangLineSearchResult(
        arg_4=arg_18,
        arg_19=arg_19 | arg_23,
        func_evals=arg_16 + arg_21,
        iterations=tf.convert_to_tensor(value=0),
        left=arg_13,
        right=hzl.val_where(arg_18, arg_13, arg_22))

    def _apply_bracket_and_search():
      """Bracketing and searching to do for valid inputs."""
      return _bracket_and_search(
          arg_0, arg_24, arg_15, arg_11,
          arg_6, arg_7, arg_8,
          arg_9)

    arg_25 = ~arg_24.failed & ~arg_24.converged
    return prefer_static.cond(
        tf.reduce_any(input_tensor=arg_25),
        _apply_bracket_and_search,
        lambda: arg_24)


def Func(arg_0,
                   arg_1,
                   arg_2,
                   arg_3):
  """Shrinks the input step size until the value and grad become finite."""
  # The maximum iterations permitted are determined as the number of halvings
  # it takes to reduce 1 to 0 in the given dtype.
  arg_4 = np.ceil(-np.log2(_machine_eps(arg_1.x.dtype)))

  def _cond(arg_5, arg_6, arg_7):
    del arg_6  # Unused.
    return (arg_5 < arg_4) & tf.reduce_any(input_tensor=arg_7)

  def _body(arg_5, arg_6, arg_7):
    arg_8 = tf.where(arg_7, arg_6.x * arg_3, arg_6.x)
    arg_9 = arg_0(arg_8)
    arg_10 = arg_7 & ~hzl.is_finite(arg_9)
    return (arg_5 + 1, arg_9, arg_10)

  arg_7 = arg_2 & ~hzl.is_finite(arg_1)
  return tf.while_loop(
      cond=_cond, body=_body, loop_vars=(0, arg_1, arg_7))


def Func(
    arg_0,
    arg_1,
    arg_2,
    arg_3,
    arg_4,
    arg_5,
    arg_6,
    arg_7):
  """The main loop of line search after the minimum has been bracketed.

  Args:
    value_and_gradients_function: A Python callable that accepts a real scalar
      tensor and returns a namedtuple with the fields 'x', 'f', and 'df' that
      correspond to scalar tensors of real dtype containing the point at which
      the function was evaluated, the value of the function, and its
      derivative at that point. The other namedtuple fields, if present,
      should be tensors or sequences (possibly nested) of tensors.
      In usual optimization application, this function would be generated by
      projecting the multivariate objective function along some specific
      direction. The direction is determined by some other procedure but should
      be a descent direction (i.e. the derivative of the projected univariate
      function must be negative at 0.).
      Alternatively, the function may represent the batching of `n` such line
      functions (e.g. projecting a single multivariate objective function along
      `n` distinct directions at once) accepting n points as input, i.e. a
      tensor of shape [n], and the fields 'x', 'f' and 'df' in the returned
      namedtuple should each be a tensor of shape [n], with the corresponding
      input points, function values, and derivatives at those input points.
    search_interval: Instance of `HagerZhangLineSearchResults` containing
      the current line search interval.
    val_0: A namedtuple as returned by value_and_gradients_function evaluated
      at `0.`. The gradient must be negative (i.e. must be a descent direction).
    f_lim: Scalar `Tensor` of float dtype.
    max_iterations: Positive scalar `Tensor` of integral dtype. The maximum
      number of iterations to perform in the line search. The number of
      iterations used to bracket the minimum are also counted against this
      parameter.
    sufficient_decrease_param: Positive scalar `Tensor` of real dtype.
      Bounded above by the curvature param. Corresponds to `delta` in the
      terminology of [Hager and Zhang (2006)][2].
    curvature_param: Positive scalar `Tensor` of real dtype. Bounded above
      by `1.`. Corresponds to 'sigma' in the terminology of
      [Hager and Zhang (2006)][2].
    shrinkage_param: Scalar positive Tensor of real dtype. Must be less than
      `1.`. Corresponds to the parameter `gamma` in [Hager and Zhang (2006)][2].

  Returns:
    A namedtuple containing the following fields.
      converged: Boolean `Tensor` of shape [n]. Whether a point satisfying
        Wolfe/Approx wolfe was found.
      failed: Boolean `Tensor` of shape [n]. Whether line search failed e.g.
        if either the objective function or the gradient are not finite at
        an evaluation point.
      iterations: Scalar int32 `Tensor`. Number of line search iterations made.
      func_evals: Scalar int32 `Tensor`. Number of function evaluations made.
      left: A namedtuple, as returned by value_and_gradients_function,
        of the left end point of the updated bracketing interval.
      right: A namedtuple, as returned by value_and_gradients_function,
        of the right end point of the updated bracketing interval.
  """

  def _loop_cond(arg_8):
    """Loop condition."""
    arg_9 = ~(arg_8.converged | arg_8.failed)
    return (arg_8.iterations <
            arg_4) & tf.reduce_any(input_tensor=arg_9)

  def _loop_body(arg_8):
    """The loop body."""
    arg_10 = hzl.secant2(
        arg_0, arg_2, arg_8, arg_3,
        arg_5, arg_6)
    arg_11 = HagerZhangLineSearchResult(
        converged=arg_10.converged,
        failed=arg_10.failed,
        iterations=arg_8.iterations + 1,
        func_evals=arg_10.num_evals,
        left=arg_10.left,
        right=arg_10.right)

    arg_12 = ~(arg_11.converged | arg_11.failed)

    def _do_check_shrinkage():
      """Check if interval has shrinked enough."""
      arg_13 = arg_8.right.x - arg_8.left.x
      arg_14 = arg_11.right.x - arg_11.left.x
      arg_15 = arg_14 < arg_13 * arg_7
      arg_16 = (
          _very_close(arg_8.left.f, arg_8.right.f) &
          _very_close(arg_11.left.f, arg_11.right.f))

      arg_17 = (
          arg_12 & arg_15 & arg_16)
      arg_18 = arg_12 & ~arg_15

      arg_19 = arg_11._replace(
          converged=arg_11.converged | arg_17)

      def _apply_inner_bisect():
        return _line_search_inner_bisection(
            arg_0, arg_19,
            arg_18, arg_3)

      return prefer_static.cond(
          tf.reduce_any(input_tensor=arg_18),
          _apply_inner_bisect,
          lambda: arg_19)

    arg_20 = prefer_static.cond(
        tf.reduce_any(input_tensor=arg_12),
        _do_check_shrinkage,
        lambda: arg_11)

    arg_21 = (
        ~arg_20.failed & _very_close(arg_20.left.x, arg_20.right.x))
    return [arg_20._replace(converged=arg_20.converged | arg_21)]

  return tf.while_loop(
      cond=_loop_cond,
      body=_loop_body,
      loop_vars=[arg_1],
      parallel_iterations=1)[0]


def Func(
    arg_0,
    arg_1,
    arg_2,
    arg_3):
  """Performs bisection and updates the interval."""
  arg_4 = (arg_1.left.x + arg_1.right.x) / 2
  arg_5 = arg_0(arg_4)
  arg_6 = hzl.is_finite(arg_5)

  arg_7 = arg_2 & arg_6
  arg_8 = arg_2 & ~arg_6
  arg_9 = arg_1._replace(
      failed=arg_1.failed | arg_8,
      func_evals=arg_1.func_evals + 1)

  def _apply_update():
    arg_10 = hzl.update(
        arg_0, arg_9.left, arg_9.right,
        arg_5, arg_3, arg_2=arg_7)
    return HagerZhangLineSearchResult(
        converged=arg_9.converged,
        failed=arg_9.failed | arg_10.failed,
        iterations=arg_9.iterations + arg_10.iteration,
        func_evals=arg_9.func_evals + arg_10.num_evals,
        left=arg_10.left,
        right=arg_10.right)

  return prefer_static.cond(
      tf.reduce_any(input_tensor=arg_7),
      _apply_update,
      lambda: arg_9)


def Func(arg_0, arg_1):
  """Wrapper for tf.Print which supports lists and namedtuples for printing."""
  arg_2 = []
  for arg_3 in arg_1:
    # Checks if it is a namedtuple.
    if hasattr(arg_3, '_fields'):
      for arg_4 in arg_3._fields:
        arg_2.extend([arg_4, _to_str(getattr(arg_3, arg_4))])
      continue
    if isinstance(arg_3, (list, tuple)):
      for arg_5 in arg_3:
        arg_2.append(_to_str(arg_5))
      continue
    arg_2.append(_to_str(arg_3))
  return tf.compat.v1.Print(arg_0, arg_2)


def Func(
    arg_0, arg_1, arg_2,
    arg_3=False, arg_4=None):
  """Use Gauss-Hermite quadrature to form quadrature on `K - 1` simplex.

  A `SoftmaxNormal` random variable `Y` may be generated via

  ```
  Y = SoftmaxCentered(X),
  X = Normal(normal_loc, normal_scale)
  ```

  Note: for a given `quadrature_size`, this method is generally less accurate
  than `quadrature_scheme_softmaxnormal_quantiles`.

  Args:
    normal_loc: `float`-like `Tensor` with shape `[b1, ..., bB, K-1]`, B>=0.
      The location parameter of the Normal used to construct the SoftmaxNormal.
    normal_scale: `float`-like `Tensor`. Broadcastable with `normal_loc`.
      The scale parameter of the Normal used to construct the SoftmaxNormal.
    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: Shape `[b1, ..., bB, K, quadrature_size]` `Tensor` representing the
      convex combination of affine parameters for `K` components.
      `grid[..., :, n]` is the `n`-th grid point, living in the `K - 1` simplex.
    probs:  Shape `[b1, ..., bB, K, quadrature_size]` `Tensor` representing the
      associated with each grid point.
  """
  with tf.name_scope(
      arg_4 or "Func"):
    arg_0 = tf.convert_to_tensor(value=arg_0, arg_4="normal_loc")
    arg_5 = dtype_util.as_numpy_dtype(arg_0.dtype)
    arg_1 = tf.convert_to_tensor(
        value=arg_1, dtype=arg_5, arg_4="normal_scale")

    arg_1 = maybe_check_quadrature_param(
        arg_1, "normal_scale", arg_3)

    arg_6, arg_7 = np.polynomial.hermite.hermgauss(deg=arg_2)
    arg_6 = arg_6.astype(arg_5)
    arg_7 = arg_7.astype(arg_5)
    arg_7 /= np.linalg.norm(arg_7, ord=1, keepdims=True)
    arg_7 = tf.convert_to_tensor(value=arg_7, arg_4="probs", dtype=arg_5)

    arg_6 = softmax(
        -distribution_util.pad(
            (arg_0[..., tf.newaxis] +
             np.sqrt(2.) * arg_1[..., tf.newaxis] * arg_6),
            axis=-2,
            front=True),
        axis=-2)  # shape: [B, components, deg]

    return arg_6, arg_7


def Func(arg_0, arg_1, arg_2):
  """Helper which checks validity of `loc` and `scale` init args."""
  with tf.name_scope("check_" + arg_1):
    arg_3 = []
    if tensorshape_util.rank(arg_0.shape) is not None:
      if tensorshape_util.rank(arg_0.shape) == 0:
        raise ValueError("Mixing params must be a (batch of) vector; "
                         "{}.rank={} is not at least one.".format(
                             arg_1, tensorshape_util.rank(arg_0.shape)))
    elif arg_2:
      arg_3.append(
          assert_util.assert_rank_at_least(
              arg_0,
              1,
              message=("Mixing params must be a (batch of) vector; "
                       "{}.rank is not at least one.".format(arg_1))))

    # TODO(jvdillon): Remove once we support k-mixtures.
    if tensorshape_util.with_rank_at_least(arg_0.shape, 1)[-1] is not None:
      if tf.compat.dimension_value(arg_0.shape[-1]) != 1:
        raise NotImplementedError("Currently only bimixtures are supported; "
                                  "{}.shape[-1]={} is not 1.".format(
                                      arg_1,
                                      tf.compat.dimension_value(
                                          arg_0.shape[-1])))
    elif arg_2:
      arg_3.append(
          assert_util.assert_equal(
              tf.shape(input=arg_0)[-1],
              1,
              message=("Currently only bimixtures are supported; "
                       "{}.shape[-1] is not 1.".format(arg_1))))

    if arg_3:
      return distribution_util.with_dependencies(arg_3, arg_0)
    return arg_0


def Func(arg_0, arg_1):
  """Helper which interpolates between two locs."""
  if len(arg_1) != 2:
    raise NotImplementedError("Currently only bimixtures are supported; "
                              "len(scale)={} is not 2.".format(len(arg_1)))
  arg_2 = tf.compat.dimension_value(
      tensorshape_util.with_rank_at_least(arg_0.shape, 1)[-1])
  if arg_2 is None:
    raise ValueError("Num quadrature grid points must be known prior "
                     "to graph execution.")
  with tf.name_scope("Func"):
    if arg_1 is None or arg_1[0] is None and arg_1[1] is None:
      return [None]*arg_2
    # shape: [B, 1, k, deg]
    arg_3 = arg_0[..., tf.newaxis, :, :]
    arg_1 = [
        arg_4[..., tf.newaxis]  # shape: [B, e, 1]
        if arg_4 is not None else None for arg_4 in arg_1
    ]
    if arg_1[0] is None:
      arg_4 = arg_3[..., 1, :] * arg_1[1]                        # shape: [B, e, deg]
    elif arg_1[1] is None:
      arg_4 = arg_3[..., 0, :] * arg_1[0]                        # shape: [B, e, deg]
    else:
      arg_5 = arg_1[0] - arg_1[1]
      arg_4 = arg_3[..., 0, :] * arg_5 + arg_1[1]                # shape: [B, e, deg]
    return [arg_4[..., arg_6] for arg_6 in range(arg_2)]


def Func(arg_0, arg_1):
  """Helper which interpolates between two scales."""
  if len(arg_1) != 2:
    raise NotImplementedError("Currently only bimixtures are supported; "
                              "len(scale)={} is not 2.".format(len(arg_1)))
  arg_2 = tf.compat.dimension_value(
      tensorshape_util.with_rank_at_least(arg_0.shape, 1)[-1])
  if arg_2 is None:
    raise ValueError("Num quadrature grid points must be known prior "
                     "to graph execution.")
  with tf.name_scope("Func"):
    return [linop_add_lib.add_operators([
        linop_scale(arg_0[..., arg_3, arg_5], arg_4)
        for arg_3, arg_4 in enumerate(arg_1)
    ])[0] for arg_5 in range(arg_2)]


def Func(arg_0, arg_1):
  """Multiply tensor of vectors by matrices assuming values stored are logs."""

  return tf.reduce_logsumexp(input_tensor=arg_0[..., tf.newaxis] + arg_1, axis=-2)


def Func(arg_0, arg_1):
  """Tabulate log probabilities from a batch of distributions."""

  arg_2 = tf.reshape(tf.range(arg_0),
                      tf.concat([[arg_0],
                                 tf.ones_like(arg_1.batch_shape_tensor())],
                                axis=0))
  return distribution_util.move_dimension(arg_1.log_prob(arg_2), 0, -1)


def Func(arg_0, arg_1, arg_2=None):
    """Compute marginal posterior distribution for each state.

    This function computes, for each time step, the marginal
    conditional probability that the hidden Markov model was in
    each possible state given the observations that were made
    at each time step.
    So if the hidden states are `z[0],...,z[num_steps - 1]` and
    the observations are `x[0], ..., x[num_steps - 1]`, then
    this function computes `P(z[i] | x[0], ..., x[num_steps - 1])`
    for all `i` from `0` to `num_steps - 1`.

    This operation is sometimes called smoothing. It uses a form
    of the forward-backward algorithm.

    Note: the behavior of this function is undefined if the
    `observations` argument represents impossible observations
    from the model.

    Args:
      observations: A tensor representing a batch of observations
        made on the hidden Markov model.  The rightmost dimension of this tensor
        gives the steps in a sequence of observations from a single sample from
        the hidden Markov model. The size of this dimension should match the
        `num_steps` parameter of the hidden Markov model object. The other
        dimensions are the dimensions of the batch and these are broadcast with
        the hidden Markov model's parameters.
      name: Python `str` name prefixed to Ops created by this class.
        Default value: "HiddenMarkovModel".

    Returns:
      posterior_marginal: A `Categorical` distribution object representing the
        marginal probability of the hidden Markov model being in each state at
        each step. The rightmost dimension of the `Categorical` distributions
        batch will equal the `num_steps` parameter providing one marginal
        distribution for each step. The other dimensions are the dimensions
        corresponding to the batch of observations.

    Raises:
      ValueError: if rightmost dimension of `observations` does not
      have size `num_steps`.
    """

    with tf.name_scope(arg_2 or "Func"):
      with tf.control_dependencies(arg_0._runtime_assertions):
        arg_3 = tf.shape(input=arg_1)

        with arg_0._observation_shape_preconditions(arg_3):
          arg_4 = arg_3[
              :-1 - arg_0._underlying_event_rank]
          arg_5 = arg_3[
              -1 - arg_0._underlying_event_rank:]

          arg_6 = tf.broadcast_dynamic_shape(arg_4,
                                                   arg_0.batch_shape_tensor())
          arg_7 = tf.broadcast_to(arg_0._log_init,
                                     tf.concat([arg_6,
                                                [arg_0._num_states]],
                                               axis=0))
          arg_8 = arg_0._log_trans

          arg_1 = tf.broadcast_to(arg_1,
                                         tf.concat([arg_6,
                                                    arg_5],
                                                   axis=0))
          arg_9 = tf.rank(arg_1)
          arg_10 = arg_0._underlying_event_rank
          arg_1 = distribution_util.move_dimension(
              arg_1, arg_9 - arg_10 - 1, 0)
          arg_1 = tf.expand_dims(
              arg_1,
              arg_9 - arg_10)
          arg_11 = arg_0._observation_distribution.log_prob(
              arg_1)

          arg_12 = tf.zeros_like(arg_7)

          def forward_step(arg_13, arg_14):
            return _log_vector_matrix(arg_13,
                                      arg_8) + arg_14

          arg_15 = arg_7 + arg_11[0]

          arg_16 = tf.scan(forward_step, arg_11[1:],
                                      initializer=arg_15,
                                      arg_2="forward_log_probs")

          arg_16 = tf.concat([[arg_15], arg_16], axis=0)

          def backward_step(arg_13, arg_14):
            return _log_matrix_vector(arg_8,
                                      arg_14 + arg_13)

          arg_17 = tf.scan(
              backward_step,
              arg_11[1:],
              initializer=arg_12,
              reverse=True,
              arg_2="backward_log_adjoint_probs")

          arg_18 = tf.reduce_logsumexp(
              input_tensor=arg_16[-1], axis=-1)

          arg_17 = tf.concat([arg_17,
                                                  [arg_12]], axis=0)

          arg_19 = arg_16 + arg_17

          arg_20 = distribution_util.move_dimension(
              arg_19 - arg_18[..., tf.newaxis], 0, -2)

          return categorical.Categorical(logits=arg_20)


def Func(arg_0, arg_1, arg_2=None):
  """Chooses a random direction in the event space."""
  arg_3 = distributions.SeedStream(arg_2, salt='Func')
  # Chooses the random directions across each of the input components.
  arg_4 = [
      tf.random.normal(
          tf.shape(input=current_state_part), dtype=tf.float32, arg_2=arg_3())
      for current_state_part in arg_0
  ]

  # Sum squares over all of the input components. Note this takes all
  # components into account.
  arg_5 = sum(
      tf.reduce_sum(
          input_tensor=rnd_direction**2.,
          axis=tf.range(arg_1, tf.rank(rnd_direction)),
          keepdims=True) for rnd_direction in arg_4)

  # Normalizes the random direction fragments.
  arg_4 = [rnd_direction / tf.sqrt(arg_5)
                         for rnd_direction in arg_4]

  return arg_4


def Func(arg_0,
                 arg_1,
                 arg_2,
                 arg_3,
                 arg_4,
                 arg_5,
                 arg_6=None,
                 arg_7=None):
  """Applies a single iteration of slice sampling update.

  Applies hit and run style slice sampling. Chooses a uniform random direction
  on the unit sphere in the event space. Applies the one dimensional slice
  sampling update along that direction.

  Args:
    target_log_prob_fn: Python callable which takes an argument like
      `*current_state_parts` and returns its (possibly unnormalized) log-density
      under the target distribution.
    current_state_parts: Python `list` of `Tensor`s representing the current
      state(s) of the Markov chain(s). The first `independent_chain_ndims` of
      the `Tensor`(s) index different chains.
    step_sizes: Python `list` of `Tensor`s. Provides a measure of the width
      of the density. Used to find the slice bounds. Must broadcast with the
      shape of `current_state_parts`.
    max_doublings: Integer number of doublings to allow while locating the slice
      boundaries.
    current_target_log_prob: `Tensor` representing the value of
      `target_log_prob_fn(*current_state_parts)`. The only reason to specify
      this argument is to reduce TF graph size.
    batch_rank: Integer. The number of axes in the state that correspond to
      independent batches.
    seed: Python integer to seed random number generators.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., 'find_slice_bounds').

  Returns:
    proposed_state_parts: Tensor or Python list of `Tensor`s representing the
      state(s) of the Markov chain(s) at each result step. Has same shape as
      input `current_state_parts`.
    proposed_target_log_prob: `Tensor` representing the value of
      `target_log_prob_fn` at `next_state`.
    bounds_satisfied: Boolean `Tensor` of the same shape as the log density.
      True indicates whether the an interval containing the slice for that
      batch was found successfully.
    direction: `Tensor` or Python list of `Tensors`s representing the direction
      along which the slice was sampled. Has the same shape and dtype(s) as
      `current_state_parts`.
    upper_bounds: `Tensor` of batch shape and the dtype of the input state. The
      upper bounds of the slices along the sampling direction.
    lower_bounds: `Tensor` of batch shape and the dtype of the input state. The
      lower bounds of the slices along the sampling direction.
  """
  with tf.compat.v1.name_scope(arg_7, 'sample_next', [
      arg_1, arg_2, arg_3, arg_4,
      arg_5
  ]):
    # First step: Choose a random direction.
    # Direction is a list of tensors. The i'th tensor should have the same shape
    # as the i'th state part.
    arg_8 = _choose_random_direction(arg_1,
                                         arg_5=arg_5,
                                         arg_6=arg_6)

    # Interpolates the step sizes for the chosen direction.
    # Applies an ellipsoidal interpolation to compute the step direction for
    # the chosen direction. Suppose we are given step sizes for each direction.
    # Label these s_1, s_2, ... s_k. These are the step sizes to use if moving
    # in a direction parallel to one of the axes. Consider an ellipsoid which
    # intercepts the i'th axis at s_i. The step size for a direction specified
    # by the unit vector (n_1, n_2 ...n_k) is then defined as the intersection
    # of the line through this vector with this ellipsoid.
    #
    # One can show that the length of the vector from the origin to the
    # intersection point is given by:
    # 1 / sqrt(n_1^2 / s_1^2  + n_2^2 / s_2^2  + ...).
    #
    # Proof:
    # The equation of the ellipsoid is:
    # Sum_i [x_i^2 / s_i^2 ] = 1. Let n be a unit direction vector. Points
    # along the line given by n may be parameterized as alpha*n where alpha is
    # the distance along the vector. Plugging this into the equation for the
    # ellipsoid, we get:
    # alpha^2 ( n_1^2 / s_1^2 + n_2^2 / s_2^2 + ...) = 1
    # so alpha = \sqrt { \frac{1} { ( n_1^2 / s_1^2 + n_2^2 / s_2^2 + ...) } }
    arg_9 = [tf.range(arg_5, tf.rank(dirn_part))
                   for dirn_part in arg_8]

    arg_10 = [
        tf.reduce_sum(
            input_tensor=(dirn_part / arg_11)**2, axis=arg_9[i])
        for i, (arg_11, dirn_part) in enumerate(zip(arg_2, arg_8))
    ]
    arg_11 = tf.math.rsqrt(tf.add_n(arg_10))
    # Computes the rank of a tensor. Uses the static rank if possible.
    def _get_rank(arg_12):
      return (len(arg_12.shape.as_list()) if arg_12.shape.dims is not None
              else tf.rank(arg_12))
    arg_13 = [_get_rank(part) for part in arg_1]

    def _step_along_direction(arg_14):
      """Converts the scalar alpha into an n-dim vector with full state info.

      Computes x_0 + alpha * direction where x_0 is the current state and
      direction is the direction chosen above.

      Args:
        alpha: A tensor of shape equal to the batch dimensions of
          `current_state_parts`.

      Returns:
        state_parts: Tensor or Python list of `Tensor`s representing the
          state(s) of the Markov chain(s) for a given alpha and a given chosen
          direction. Has the same shape as `current_state_parts`.
      """
      arg_15 = [_right_pad(arg_14, final_rank=part_rank)
                       for part_rank in arg_13]

      arg_16 = [state_part + padded_alpha
                     * direction_part for state_part, direction_part,
                     padded_alpha in
                     zip(arg_1, arg_8, arg_15)]
      return arg_16

    def projected_target_log_prob_fn(arg_14):
      """The target log density projected along the chosen direction.

      Args:
        alpha: A tensor of shape equal to the batch dimensions of
          `current_state_parts`.

      Returns:
        Target log density evaluated at x_0 + alpha * direction where x_0 is the
        current state and direction is the direction chosen above. Has the same
        shape as `alpha`.
      """
      return arg_0(*_step_along_direction(arg_14))

    arg_17 = tf.zeros_like(arg_4,
                               dtype=arg_1[0].dtype.base_dtype)
    [
        arg_18,
        arg_19,
        arg_20,
        arg_21,
        arg_22
    ] = ssu.slice_sampler_one_dim(projected_target_log_prob_fn,
                                  x_initial=arg_17,
                                  arg_3=arg_3,
                                  arg_11=arg_11, arg_6=arg_6)
    return [
        _step_along_direction(arg_18),
        arg_19,
        arg_20,
        arg_8,
        arg_21,
        arg_22
    ]