def Func(arg_0,
                     arg_1,
                     arg_2,
                     arg_3,
                     arg_4):
  """Checks whether the Wolfe or approx Wolfe conditions are satisfied.

  The Wolfe conditions are a set of stopping criteria for an inexact line search
  algorithm. Let f(a) be the function value along the search direction and
  df(a) the derivative along the search direction evaluated a distance 'a'.
  Here 'a' is the distance along the search direction. The Wolfe conditions are:

    ```None
      f(a) <= f(0) + delta * a * df(0)   (Armijo/Sufficient decrease condition)
      df(a) >= sigma * df(0)             (Weak curvature condition)
    ```
  `delta` and `sigma` are two user supplied parameters satisfying:
   `0 < delta < sigma <= 1.`. In the following, delta is called
   `sufficient_decrease_param` and sigma is called `curvature_param`.

  On a finite precision machine, the Wolfe conditions are difficult to satisfy
  when one is close to the minimum. Hence, Hager-Zhang propose replacing
  the sufficient decrease condition with the following condition on the
  derivative in the vicinity of a minimum.

    ```None
      df(a) <= (2 * delta - 1) * df(0)  (Approx Wolfe sufficient decrease)
    ```
  This condition is only used if one is near the minimum. This is tested using

    ```None
      f(a) <= f(0) + epsilon * |f(0)|
    ```
  The following function checks both the Wolfe and approx Wolfe conditions.
  Here, `epsilon` is a small positive constant. In the following, the argument
  `f_lim` corresponds to the product: epsilon * |f(0)|.

  Args:
    val_0: A namedtuple, as returned by value_and_gradients_function
      evaluated at 0.
    val_c: A namedtuple, as returned by value_and_gradients_function
      evaluated at the point to be tested.
    f_lim: Scalar `Tensor` of real dtype. The function value threshold for
      the approximate Wolfe conditions to be checked.
    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 Zhang (2005)][1].

  Returns:
    is_satisfied: A scalar boolean `Tensor` which is True if either the
      Wolfe or approximate Wolfe conditions are satisfied.
  """
  arg_5 = (arg_3 * arg_0.df >=
                          (arg_1.f - arg_0.f) / arg_1.x)
  arg_6 = arg_1.df >= arg_4 * arg_0.df
  arg_7 = arg_5 & arg_6
  arg_8 = arg_1.f <= arg_2
  arg_9 = ((2 * arg_3 - 1) * arg_0.df
                           >= arg_1.df)
  arg_10 = arg_8 & arg_9 & arg_6
  arg_11 = arg_7 | arg_10
  return arg_11


def Func(arg_0, arg_1):
  """Returns the secant interpolation for the minimum.

  The secant method is a technique for finding roots of nonlinear functions.
  When finding the minimum, one applies the secant method to the derivative
  of the function.
  For an arbitrary function and a bounding interval, the secant approximation
  can produce the next point which is outside the bounding interval. However,
  with the assumption of opposite slope condtion on the interval [a,b] the new
  point c is always bracketed by [a,b]. Note that by assumption,
  f'(a) < 0 and f'(b) > 0.
  Hence c is a weighted average of a and b and thus always in [a, b].

  Args:
    val_a: A namedtuple with the left end point, function value and derivative,
      of the current interval (i.e. a).
    val_b: A namedtuple with the right end point, function value and derivative,
      of the current interval (i.e. b).

  Returns:
    approx_minimum: A scalar real `Tensor`. An approximation to the point
      at which the derivative vanishes.
  """
  return (arg_0.x * arg_1.df - arg_1.x * arg_0.df) / (arg_1.df - arg_0.df)


def Func(arg_0,
                                        arg_1=0.75,
                                        arg_2=0.01,
                                        arg_3=0.01,
                                        arg_4=None):
  """Create a function implementing a step-size update policy.

  The simple policy increases or decreases the `step_size_var` based on the
  average of `exp(minimum(0., log_accept_ratio))`. It is based on
  [Section 4.2 of Andrieu and Thoms (2008)](
  https://people.eecs.berkeley.edu/~jordan/sail/readings/andrieu-thoms.pdf).

  The `num_adaptation_steps` argument is set independently of any burnin
  for the overall chain. In general, adaptation prevents the chain from
  reaching a stationary distribution, so obtaining consistent samples requires
  `num_adaptation_steps` be set to a value [somewhat smaller](
  http://andrewgelman.com/2017/12/15/burn-vs-warm-iterative-simulation-algorithms/#comment-627745)
  than the number of burnin steps. However, it may sometimes be helpful to set
  `num_adaptation_steps` to a larger value during development in order to
  inspect the behavior of the chain during adaptation.

  Args:
    num_adaptation_steps: Scalar `int` `Tensor` number of initial steps to
      during which to adjust the step size. This may be greater, less than, or
      equal to the number of burnin steps. If `None`, the step size is adapted
      on every step (note this breaks stationarity of the chain!).
    target_rate: Scalar `Tensor` representing desired `accept_ratio`.
      Default value: `0.75` (i.e., [center of asymptotically optimal
      rate](https://arxiv.org/abs/1411.6669)).
    decrement_multiplier: `Tensor` representing amount to downscale current
      `step_size`.
      Default value: `0.01`.
    increment_multiplier: `Tensor` representing amount to upscale current
      `step_size`.
      Default value: `0.01`.
    step_counter: Scalar `int` `Variable` specifying the current step. The step
      size is adapted iff `step_counter < num_adaptation_steps`.
      Default value: if `None`, an internal variable
        `step_size_adaptation_step_counter` is created and initialized to `-1`.

  Returns:
    step_size_simple_update_fn: Callable that takes args
      `step_size_var, kernel_results` and returns updated step size(s).
  """
  if arg_4 is None and arg_0 is not None:
    arg_4 = tf.compat.v1.get_variable(
        name='step_size_adaptation_step_counter',
        initializer=np.array(-1, dtype=np.int32),
        # Specify the dtype for variable sharing to work correctly
        # (b/120599991).
        dtype=tf.int32,
        trainable=False,
        use_resource=True)

  def step_size_simple_update_fn(arg_5, arg_6):
    """Updates (list of) `step_size` using a standard adaptive MCMC procedure.

    Args:
      step_size_var: (List of) `tf.Variable`s representing the per `state_part`
        HMC `step_size`.
      kernel_results: `collections.namedtuple` containing `Tensor`s
        representing values from most recent call to `one_step`.

    Returns:
      step_size_assign: (List of) `Tensor`(s) representing updated
        `step_size_var`(s).
    """

    if arg_6 is None:
      if mcmc_util.is_list_like(arg_5):
        return [tf.identity(arg_7) for arg_7 in arg_5]
      return tf.identity(arg_5)
    arg_8 = tf.math.log(
        tf.cast(
            tf.size(input=arg_6.log_accept_ratio),
            arg_6.log_accept_ratio.dtype))
    arg_9 = tf.reduce_logsumexp(
        input_tensor=tf.minimum(arg_6.log_accept_ratio, 0.)) - arg_8
    arg_10 = tf.where(
        arg_9 < tf.cast(
            tf.math.log(arg_1), arg_9.dtype),
        -arg_2 / (1. + arg_2),
        arg_3)

    def build_assign_op():
      if mcmc_util.is_list_like(arg_5):
        return [
            arg_7.assign_add(arg_7 * tf.cast(arg_10, arg_7.dtype))
            for arg_7 in arg_5
        ]
      return arg_5.assign_add(
          arg_5 * tf.cast(arg_10, arg_5.dtype))

    if arg_0 is None:
      return build_assign_op()
    else:
      with tf.control_dependencies([arg_4.assign_add(1)]):
        return tf.cond(
            pred=arg_4 < arg_0,
            true_fn=build_assign_op,
            false_fn=lambda: arg_5)

  return step_size_simple_update_fn


def Func(arg_0, arg_1):
    """Creates initial `previous_kernel_results` using a supplied `state`."""
    arg_2 = arg_0._impl.Func(arg_1)
    if arg_0.step_size_update_fn is not None:
      arg_3 = arg_0.step_size_update_fn(arg_0.step_size, None)  # pylint: disable=not-callable
      arg_2 = arg_2._replace(
          extra=HamiltonianMonteCarloExtraKernelResults(
              arg_3=arg_3))
    return arg_2


def Func(arg_0,
                    arg_1=10,
                    arg_2=-9.0,
                    arg_3=0.1,
                    arg_4=0.2):
  """Constructs a ResNet18 model.

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

  Returns:
    tf.keras.Model.
  """

  arg_5 = [64, 128, 256, 512]
  arg_6 = [3, 3, 3, 3]
  arg_7 = [1, 2, 2, 2]

  def _untransformed_scale_constraint(arg_8):
    return tf.clip_by_value(arg_8, -1000,
                            tf.math.log(arg_4))

  arg_9 = tfp.layers.default_mean_field_normal_fn(
      untransformed_scale_initializer=tf.compat.v1.initializers.random_normal(
          mean=arg_2,
          stddev=arg_3),
      untransformed_scale_constraint=_untransformed_scale_constraint)

  arg_10 = tf.keras.layers.Input(shape=arg_0, dtype='float32')
  arg_11 = tfp.layers.Convolution2DFlipout(
      64,
      3,
      arg_7=1,
      padding='same',
      arg_9=arg_9)(arg_10)

  for arg_12 in range(len(arg_6)):
    arg_11 = _resnet_block(
        arg_11,
        arg_5[arg_12],
        arg_6[arg_12],
        arg_7[arg_12],
        arg_9)

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

  arg_11 = tfp.layers.DenseFlipout(
      arg_1,
      arg_9=arg_9)(arg_11)

  arg_13 = tf.keras.Model(inputs=arg_10, outputs=arg_11, name='resnet18')
  return arg_13


def Func(arg_0, arg_1, arg_2):
  """Create the encoder function.

  Args:
    activation: Activation function to use.
    num_topics: The number of topics.
    layer_sizes: The number of hidden units per layer in the encoder.

  Returns:
    encoder: A `callable` mapping a bag-of-words `Tensor` to a
      `tfd.Distribution` instance over topics.
  """
  arg_3 = tf.keras.Sequential()
  for arg_4 in arg_2:
    arg_3.add(
        tf.keras.layers.Dense(
            arg_4,
            arg_0=arg_0,
            kernel_initializer=tf.compat.v1.glorot_normal_initializer()))
  arg_3.add(
      tf.keras.layers.Dense(
          arg_1,
          arg_0=tf.nn.softplus,
          kernel_initializer=tf.compat.v1.glorot_normal_initializer()))

  def encoder(arg_5):
    arg_6 = _clip_dirichlet_parameters(arg_3(arg_5))
    return tfd.Dirichlet(concentration=arg_6,
                         name="topics_posterior")

  return encoder


def Func(arg_0, arg_1):
  """Create the decoder function.

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

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

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

  return decoder, arg_3


def Func(
    arg_0,
    arg_1,
    arg_2=None,
    arg_3=None,
    arg_4=0,
    arg_5=0,
    arg_6=lambda arg_1, arg_7: arg_7,
    arg_8=False,
    arg_9=10,
    arg_10=None,
):
  """Implements Markov chain Monte Carlo via repeated `TransitionKernel` steps.

  This function samples from an Markov chain at `current_state` and whose
  stationary distribution is governed by the supplied `TransitionKernel`
  instance (`kernel`).

  This function can sample from multiple chains, in parallel. (Whether or not
  there are multiple chains is dictated by the `kernel`.)

  The `current_state` can be represented as a single `Tensor` or a `list` of
  `Tensors` which collectively represent the current state.

  Since MCMC states are correlated, it is sometimes desirable to produce
  additional intermediate states, and then discard them, ending up with a set of
  states with decreased autocorrelation.  See [Owen (2017)][1]. Such "thinning"
  is made possible by setting `num_steps_between_results > 0`. The chain then
  takes `num_steps_between_results` extra steps between the steps that make it
  into the results. The extra steps are never materialized (in calls to
  `sess.run`), and thus do not increase memory requirements.

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

  In addition to returning the chain state, this function supports tracing of
  auxiliary variables used by the kernel. The traced values are selected by
  specifying `trace_fn`. By default, all kernel results are traced but in the
  future the default will be changed to no results being traced, so plan
  accordingly. See below for some examples of this feature.

  Args:
    num_results: Integer number of Markov chain draws.
    current_state: `Tensor` or Python `list` of `Tensor`s representing the
      current state(s) of the Markov chain(s).
    previous_kernel_results: A `Tensor` or a nested collection of `Tensor`s
      representing internal calculations made within the previous call to this
      function (or as returned by `bootstrap_results`).
    kernel: An instance of `tfp.mcmc.TransitionKernel` which implements one step
      of the Markov chain.
    num_burnin_steps: Integer number of chain steps to take before starting to
      collect results.
      Default value: 0 (i.e., no burn-in).
    num_steps_between_results: Integer number of chain steps between collecting
      a result. Only one out of every `num_steps_between_samples + 1` steps is
      included in the returned results.  The number of returned chain states is
      still equal to `num_results`.  Default value: 0 (i.e., no thinning).
    trace_fn: A callable that takes in the current chain state and the previous
      kernel results and return a `Tensor` or a nested collection of `Tensor`s
      that is then traced along with the chain state.
    return_final_kernel_results: If `True`, then the final kernel results are
      returned alongside the chain state and the trace specified by the
      `trace_fn`.
    parallel_iterations: The number of iterations allowed to run in parallel. It
      must be a positive integer. See `tf.while_loop` for more details.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., "mcmc_Func").

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

  #### Examples

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

  I.e.,

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

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

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

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

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

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

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

  I.e.,

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

  where `F` denotes factors.

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

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

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

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

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

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

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

  # Sample from Hamiltonian Monte Carlo Markov Chain.

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

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

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

  ##### Custom tracing functions.

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

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

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

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

  def trace_everything(states, previous_kernel_results):
    return previous_kernel_results

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

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

  #### References

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

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

    def _trace_scan_fn(arg_12, arg_13):
      arg_14, arg_15 = mcmc_util.smart_for_loop(
          loop_num_iter=arg_13,
          body_fn=arg_3.one_step,
          initial_loop_vars=list(arg_12),
          arg_9=arg_9)
      return arg_14, arg_15

    (arg_16, arg_17), (arg_18, arg_19) = mcmc_util.trace_scan(
        loop_fn=_trace_scan_fn,
        initial_state=(arg_1, arg_2),
        elems=tf.one_hot(
            indices=0,
            depth=arg_0,
            on_value=1 + arg_4,
            off_value=1 + arg_5,
            dtype=tf.int32),
        # pylint: disable=g-long-lambda
        arg_6=lambda arg_12: (arg_12[0],
                                            arg_6(*arg_12)),
        # pylint: enable=g-long-lambda
        arg_9=arg_9)

    if arg_8:
      return CheckpointableStatesAndTrace(
          arg_18=arg_18,
          arg_19=arg_19,
          arg_17=arg_17)
    else:
      if arg_11:
        return arg_18
      else:
        return StatesAndTrace(arg_18=arg_18, arg_19=arg_19)


def Func(arg_0, arg_1, arg_2, arg_3):
  """A multi-layered topic model over a documents-by-terms matrix."""
  arg_4 = ed.Gamma(0.1, 0.3, sample_shape=[arg_2[2], arg_2[1]], name="w2")
  arg_5 = ed.Gamma(0.1, 0.3, sample_shape=[arg_2[1], arg_2[0]], name="w1")
  arg_6 = ed.Gamma(0.1, 0.3, sample_shape=[arg_2[0], arg_1], name="w0")

  arg_7 = ed.Gamma(0.1, 0.1, sample_shape=[arg_0, arg_2[2]], name="z2")
  arg_8 = ed.Gamma(arg_3, arg_3 / tf.matmul(arg_7, arg_4), name="z1")
  arg_9 = ed.Gamma(arg_3, arg_3 / tf.matmul(arg_8, arg_5), name="z0")
  arg_10 = ed.Poisson(tf.matmul(arg_9, arg_6), name="x")
  return arg_10


def Func(arg_0, arg_1=1e-3, arg_2=1e-5, arg_3=None):
  """Learnable Gamma via concentration and scale parameterization."""
  with tf.compat.v1.variable_scope(None, default_name="Func"):
    arg_4 = tf.compat.v1.get_variable(
        "unconstrained_concentration",
        arg_0,
        initializer=tf.compat.v1.initializers.random_normal(
            mean=0.5, stddev=0.1))
    arg_5 = tf.compat.v1.get_variable(
        "unconstrained_scale",
        arg_0,
        initializer=tf.compat.v1.initializers.random_normal(stddev=0.1))
    arg_6 = tf.maximum(tf.nn.softplus(arg_4),
                               arg_1)
    arg_7 = tf.maximum(1. / tf.nn.softplus(arg_5), 1. / arg_2)
    arg_8 = ed.Gamma(arg_6=arg_6, arg_7=arg_7, arg_3=arg_3)
    return arg_8


def Func(arg_0, arg_1):
  """Get the KL function registered for classes a and b."""
  arg_2 = tf_inspect.getmro(arg_0)
  arg_3 = tf_inspect.getmro(arg_1)
  arg_4 = None
  arg_5 = None
  for arg_6, arg_7 in enumerate(arg_2):
    for arg_8, arg_9 in enumerate(arg_3):
      arg_10 = arg_6 + arg_8
      arg_11 = _DIVERGENCES.get((arg_7, arg_9), None)
      if not arg_5 or (arg_11 and arg_10 < arg_4):
        arg_4 = arg_10
        arg_5 = arg_11
  return arg_5


def Func(arg_0):
  """Returns an image tensor."""
  arg_1 = tf.io.read_file(arg_0)
  arg_2 = tf.image.decode_image(arg_1, channels=CHANNELS)
  arg_2 = tf.image.convert_image_dtype(arg_2, tf.float32)
  return arg_2


def Func(arg_0, arg_1, arg_2, arg_3):
  """Creates a character sprite from a set of attribute sprites."""
  arg_4 = arg_0.dtype
  arg_5 = tf.cast(arg_1[..., -1:] <= 0, arg_4)
  arg_6 = tf.cast(arg_2[..., -1:] <= 0, arg_4)
  arg_7 = tf.cast(arg_3[..., -1:] <= 0, arg_4)
  arg_8 = (arg_0 * arg_5) + arg_1
  arg_8 = (arg_8 * arg_6) + arg_2
  arg_8 = (arg_8 * arg_7) + arg_3
  return arg_8


def Func(arg_0, arg_1, arg_2, arg_3=8):
  """Creates a random sequence."""
  arg_4 = tf.random.uniform([], maxval=arg_1[1], dtype=tf.int32)
  return create_seq(arg_0, arg_1, arg_2, arg_3, arg_4)


def Func(arg_0, arg_1, arg_2, arg_3=3,
                           arg_4=8, arg_5=False, arg_6=False):
  """Creates a tf.data pipeline for the sprites dataset.

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

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

  arg_9 = [action.name for action in arg_1]
  arg_10 = [(action.start_row, action.frames) for action in arg_1]

  arg_11 = [direction.row_offset for direction in arg_2]

  arg_12 = tf.data.Dataset.from_tensor_slices(arg_0)
  arg_13 = tf.data.Dataset.from_tensor_slices(arg_9).repeat()
  arg_14 = tf.data.Dataset.from_tensor_slices(arg_10).repeat()
  arg_15 = tf.data.Dataset.from_tensor_slices(arg_11).repeat()

  if arg_5:
    arg_12 = arg_12.shuffle(len(arg_0))

  arg_16 = tf.data.Dataset.zip((arg_12, arg_13, arg_14, arg_15))

  arg_17 = tf.contrib.lookup.index_table_from_tensor(sorted(SKIN_COLORS))
  arg_18 = tf.contrib.lookup.index_table_from_tensor(sorted(HAIRSTYLES))
  arg_19 = tf.contrib.lookup.index_table_from_tensor(sorted(TOPS))
  arg_20 = tf.contrib.lookup.index_table_from_tensor(sorted(PANTS))
  arg_21 = tf.contrib.lookup.index_table_from_tensor(sorted(arg_9))

  def process_example(arg_22, arg_23, arg_24, arg_25):
    """Processes a dataset row."""
    arg_26 = arg_22[0]
    arg_27 = arg_22[1]
    arg_28 = arg_22[2]
    arg_29 = arg_22[3]

    if arg_6:
      arg_30 = arg_7
    else:
      arg_31 = read_image(arg_8 + os.sep + arg_26)
      arg_32 = read_image(arg_8 + os.sep + arg_27)
      arg_33 = read_image(arg_8 + os.sep + arg_28)
      arg_34 = read_image(arg_8 + os.sep + arg_29)
      arg_30 = create_character(arg_31, arg_32, arg_33, arg_34)

    if arg_5:
      arg_35 = create_random_seq(arg_30, arg_24, arg_25, arg_4)
    else:
      arg_35 = create_seq(arg_30, arg_24, arg_25, arg_4)
    arg_35 = arg_35[..., :arg_3]  # limit output channels

    arg_36 = arg_17.lookup(arg_26)
    arg_37 = arg_18.lookup(arg_27)
    arg_38 = arg_19.lookup(arg_28)
    arg_39 = arg_20.lookup(arg_29)
    arg_40 = arg_21.lookup(arg_23)

    return (arg_35, arg_36, arg_37, arg_38, arg_39, arg_40,
            arg_26, arg_27, arg_28, arg_29, arg_23)

  arg_16 = arg_16.map(process_example)
  return arg_16


def Func(arg_0, arg_1, arg_2):
  """Checks that `distributions` satisfies all assumptions."""
  arg_3 = []

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

  if arg_1 is None:
    arg_4 = [
        dtype_util.base_dtype(arg_6.dtype)
        for arg_6 in arg_0
        if arg_6.dtype is not None
    ]
    if arg_4[1:] != arg_4[:-1]:
      raise TypeError('Distributions must have same dtype; found: {}.'.format(
          set(dtype_util.name(arg_5) for arg_5 in arg_4)))

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

  arg_7 = [arg_6.batch_shape for arg_6 in arg_0]
  if all(tensorshape_util.is_fully_defined(arg_8) for arg_8 in arg_7):
    if arg_7[1:] != arg_7[:-1]:
      raise ValueError('Distributions must have the same `batch_shape`; '
                       'found: {}.'.format(arg_7))
  elif arg_2:
    arg_7 = [
        tensorshape_util.as_list(arg_6.batch_shape)  # pylint: disable=g-complex-comprehension
        if tensorshape_util.is_fully_defined(arg_6.batch_shape) else
        arg_6.batch_shape_tensor() for arg_6 in arg_0
    ]
    arg_3.extend(
        assert_util.assert_equal(  # pylint: disable=g-complex-comprehension
            arg_9, arg_10,
            message='Distribution `batch_shape`s must be identical.')
        for arg_9, arg_10 in zip(arg_7[1:], arg_7[:-1]))

  return arg_3


def Func(arg_0,
                   arg_1=None,
                   arg_2=None,
                   arg_3=None,
                   arg_4=None,
                   arg_5=arg_6.int32,
                   arg_8=None):
  """Counts the number of occurrences of each value in an integer array `arr`.

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

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

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

  Returns:
    A vector with the same dtype as `weights` or the given `dtype`. The bin
    values.
  """
  with arg_6.compat.v1.name_scope(
      arg_8, 'Func', values=[arg_0, arg_1, arg_2, arg_3,
                                      arg_4]):
    if arg_4 is None:
      return arg_6.math.bincount(
          arg_0,
          arg_1=arg_1,
          arg_2=arg_2,
          arg_3=arg_3,
          arg_5=arg_5)

    arg_0 = arg_6.convert_to_tensor(value=arg_0, arg_5=arg_6.int32, arg_8='arr')
    arg_9 = _get_static_ndims(arg_0, expect_static=True)

    arg_4 = _make_static_axis_non_negative_list(arg_4, arg_9)

    # ~axis from docstring.  Dims in arr that are not in axis.
    arg_10 = sorted(set(range(arg_9)).difference(arg_4))

    # If we're reducing over everything, just use standard bincount.
    if not arg_10:
      return arg_6.math.bincount(
          arg_0,
          arg_1=arg_1,
          arg_2=arg_2,
          arg_3=arg_3,
          arg_5=arg_5)

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

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

      def one_bincount(arg_12):
        return arg_6.math.bincount(
            arg_12,
            arg_1=None,
            arg_2=arg_2,
            arg_3=arg_3,
            arg_5=arg_5)

      arg_13 = arg_6.map_fn(one_bincount, elems=arg_11, arg_5=arg_5)
    else:
      arg_1 = arg_6.convert_to_tensor(value=arg_1, arg_8='weights')
      _get_static_ndims(arg_1, expect_static=True, expect_ndims=arg_9)
      arg_14 = _move_dims_to_flat_end(
          arg_1, arg_10, arg_9, right_end=False)

      def one_bincount(arg_15):
        arg_12, arg_16 = arg_15
        return arg_6.math.bincount(
            arg_12,
            arg_1=arg_16,
            arg_2=arg_2,
            arg_3=arg_3,
            arg_5=arg_5)

      arg_13 = arg_6.map_fn(
          one_bincount, elems=[arg_11, arg_14], arg_5=arg_1.dtype)

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

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

    # not_axis_shape = arr.shape[~axis]
    arg_18 = arg_6.gather(arg_6.shape(input=arg_0), indices=arg_10)

    # The first index of flat_counts_t indexes bins 0,..,K-1, the rest are ~axis
    arg_19 = arg_6.concat([[-1], arg_18], arg_4=0)

    return arg_6.reshape(arg_17, arg_19)


def Func(arg_0,
              arg_1,
              arg_2=False,
              arg_3=False,
              arg_4=None,
              arg_5=None):
  """Bin values into discrete intervals.

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

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

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

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

  #### Examples

  Cut a `1-D` array

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

  Cut `x` into its deciles

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

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

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

    arg_7 = arg_1.shape.ndims == 1 and arg_0.shape.ndims > 1

    if arg_7:
      arg_8 = tf.shape(input=arg_0)
      arg_0 = tf.reshape(arg_0, [-1])

    if arg_4 is None:
      arg_4 = arg_6
    arg_4 = tf.as_dtype(arg_4)

    # Move first dims into the rightmost.
    arg_9 = distribution_util.rotate_transpose(arg_0, shift=-1)
    arg_10 = distribution_util.rotate_transpose(arg_1, shift=-1)

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

    # In above example, we want [0, 0, 1, 1], so correct this here.
    arg_14 = tf.clip_by_value(arg_13 - 1, tf.cast(0, arg_4),
                            tf.cast(tf.shape(input=arg_1)[0] - 2, arg_4))

    if not arg_2:
      arg_15 = np.nan if arg_4.is_floating else -1
      arg_14 = tf.where(arg_0 < tf.expand_dims(arg_1[0], 0),
                      tf.fill(tf.shape(input=arg_0), tf.cast(arg_15, arg_4)),
                      arg_14)

    if not arg_3:
      arg_16 = np.nan if arg_4.is_floating else tf.shape(input=arg_1)[0] - 1
      arg_14 = tf.where(arg_0 > tf.expand_dims(arg_1[-1], 0),
                      tf.fill(tf.shape(input=arg_0), tf.cast(arg_16, arg_4)), arg_14)

    if arg_7:
      arg_14 = tf.reshape(arg_14, arg_8)

    return arg_14


def Func(arg_0,
              arg_1,
              arg_2=None,
              arg_3=False,
              arg_4=False,
              arg_5=None,
              arg_6=None):
  """Count how often `x` falls in intervals defined by `edges`.

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

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

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

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

  #### Examples

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

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

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

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

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

    # Tensor conversions.
    arg_7 = dtype_util.common_dtype([arg_0, arg_1], preferred_dtype=tf.float32)

    arg_0 = tf.convert_to_tensor(value=arg_0, arg_6='x', arg_5=arg_7)
    arg_1 = tf.convert_to_tensor(value=arg_1, arg_6='edges', arg_5=arg_7)

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

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

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


def Func(arg_0,
                      arg_1=False,
                      arg_2=None,
                      arg_3=None,
                      arg_4=None):
  """Get static number of dimensions and assert that some expectations are met.

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

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

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

  Returns:
    ndims:  A Python integer.

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

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

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

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

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

  return arg_5


def Func(arg_0, arg_1):
  """Insert the dims in `axis` back as singletons after being removed.

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

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


def Func(arg_0, arg_1):
  """Convert possibly negatively indexed axis to non-negative list of ints.

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

  Returns:
    A list of non-negative Python integers.

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

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

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

  return list(int(arg_3) for arg_3 in arg_0)


def Func(arg_0):
  """Use `top_k` to sort a `Tensor` along the last dimension."""
  arg_1, arg_2 = tf.nn.top_k(arg_0, k=tf.shape(input=arg_0)[-1])
  arg_1.set_shape(arg_0.shape)
  return arg_1


def Func(arg_0,
                                        arg_1,
                                        arg_2,
                                        arg_3=0):
    """Build an ordered list of Distribution instances for component models.

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

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

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

      # List the model parameters in canonical order
      arg_4 = arg_0._canonicalize_param_vals_as_map(arg_2)
      arg_5 = [arg_4[p.name] for p in arg_0.parameters]

      # Build SSMs for each component model. We process the components in
      # canonical order, extracting the parameters for each component from the
      # (ordered) list of parameters.
      arg_6 = arg_5[1:]
      arg_7 = []
      for arg_8 in arg_0.components:
        arg_9 = len(arg_8.parameters)
        arg_10 = arg_6[:arg_9]
        arg_6 = arg_6[arg_9:]

        arg_7.append(
            arg_8.make_state_space_model(
                arg_1,
                arg_2=arg_10,
                arg_3=arg_3))

    return arg_7


def Func(arg_0, arg_1=1., arg_2=False, arg_3=None):
  """The Amari-alpha Csiszar-function in log-space.

  A Csiszar-function is a member of,

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

  When `self_normalized = True`, the Amari-alpha Csiszar-function is:

  ```none
  f(u) = { -log(u) + (u - 1),     alpha = 0
         { u log(u) - (u - 1),    alpha = 1
         { [(u**alpha - 1) - alpha (u - 1)] / (alpha (alpha - 1)),    otherwise
  ```

  When `self_normalized = False` the `(u - 1)` terms are omitted.

  Warning: when `alpha != 0` and/or `self_normalized = True` this function makes
  non-log-space calculations and may therefore be numerically unstable for
  `|logu| >> 0`.

  For more information, see:
    A. Cichocki and S. Amari. "Families of Alpha-Beta-and GammaDivergences:
    Flexible and Robust Measures of Similarities." Entropy, vol. 12, no. 6, pp.
    1532-1568, 2010.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    alpha: `float`-like Python scalar. (See Mathematical Details for meaning.)
    self_normalized: Python `bool` indicating whether `f'(u=1)=0`. When
      `f'(u=1)=0` the implied Csiszar f-Divergence remains non-negative even
      when `p, q` are unnormalized measures.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    Func_of_u: `float`-like `Tensor` of the Csiszar-function evaluated
      at `u = exp(logu)`.

  Raises:
    TypeError: if `alpha` is `None` or a `Tensor`.
    TypeError: if `self_normalized` is `None` or a `Tensor`.
  """
  with tf.compat.v1.name_scope(arg_3, "Func", [arg_0]):
    if arg_1 is None or tf.is_tensor(arg_1):
      raise TypeError("`alpha` cannot be `None` or `Tensor` type.")
    if (arg_2 is None or tf.is_tensor(arg_2)):
      raise TypeError("`self_normalized` cannot be `None` or `Tensor` type.")

    arg_0 = tf.convert_to_tensor(value=arg_0, arg_3="logu")

    if arg_1 == 0.:
      arg_4 = -arg_0
    elif arg_1 == 1.:
      arg_4 = tf.exp(arg_0) * arg_0
    else:
      arg_4 = tf.math.expm1(arg_1 * arg_0) / (arg_1 * (arg_1 - 1.))

    if not arg_2:
      return arg_4

    if arg_1 == 0.:
      return arg_4 + tf.math.expm1(arg_0)
    elif arg_1 == 1.:
      return arg_4 - tf.math.expm1(arg_0)
    else:
      return arg_4 - tf.math.expm1(arg_0) / (arg_1 - 1.)


def Func(arg_0, arg_1=False, arg_2=None):
  """The reverse Kullback-Leibler Csiszar-function in log-space.

  A Csiszar-function is a member of,

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

  When `self_normalized = True`, the KL-reverse Csiszar-function is:

  ```none
  f(u) = -log(u) + (u - 1)
  ```

  When `self_normalized = False` the `(u - 1)` term is omitted.

  Observe that as an f-Divergence, this Csiszar-function implies:

  ```none
  D_f[p, q] = KL[q, p]
  ```

  The KL is "reverse" because in maximum likelihood we think of minimizing `q`
  as in `KL[p, q]`.

  Warning: when self_normalized = True` this function makes non-log-space
  calculations and may therefore be numerically unstable for `|logu| >> 0`.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    self_normalized: Python `bool` indicating whether `f'(u=1)=0`. When
      `f'(u=1)=0` the implied Csiszar f-Divergence remains non-negative even
      when `p, q` are unnormalized measures.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    Func_of_u: `float`-like `Tensor` of the Csiszar-function evaluated at
      `u = exp(logu)`.

  Raises:
    TypeError: if `self_normalized` is `None` or a `Tensor`.
  """

  with tf.compat.v1.name_scope(arg_2, "Func", [arg_0]):
    return amari_alpha(arg_0, alpha=0., arg_1=arg_1)


def Func(arg_0, arg_1=False, arg_2=None):
  """The Jensen-Shannon Csiszar-function in log-space.

  A Csiszar-function is a member of,

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

  When `self_normalized = True`, the Jensen-Shannon Csiszar-function is:

  ```none
  f(u) = u log(u) - (1 + u) log(1 + u) + (u + 1) log(2)
  ```

  When `self_normalized = False` the `(u + 1) log(2)` term is omitted.

  Observe that as an f-Divergence, this Csiszar-function implies:

  ```none
  D_f[p, q] = KL[p, m] + KL[q, m]
  m(x) = 0.5 p(x) + 0.5 q(x)
  ```

  In a sense, this divergence is the "reverse" of the Arithmetic-Geometric
  f-Divergence.

  This Csiszar-function induces a symmetric f-Divergence, i.e.,
  `D_f[p, q] = D_f[q, p]`.

  Warning: this function makes non-log-space calculations and may therefore be
  numerically unstable for `|logu| >> 0`.

  For more information, see:
    Lin, J. "Divergence measures based on the Shannon entropy." IEEE Trans.
    Inf. Th., 37, 145-151, 1991.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    self_normalized: Python `bool` indicating whether `f'(u=1)=0`. When
      `f'(u=1)=0` the implied Csiszar f-Divergence remains non-negative even
      when `p, q` are unnormalized measures.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    Func_of_u: `float`-like `Tensor` of the Csiszar-function
      evaluated at `u = exp(logu)`.
  """

  with tf.compat.v1.name_scope(arg_2, "Func", [arg_0]):
    arg_0 = tf.convert_to_tensor(value=arg_0, arg_2="logu")
    arg_3 = arg_0.dtype.as_numpy_dtype
    arg_4 = tf.nn.softplus(arg_0)
    if arg_1:
      arg_4 -= np.log(2).astype(arg_3)
    return tf.exp(arg_0) * arg_0 - (1. + tf.exp(arg_0)) * arg_4


def Func(arg_0, arg_1=None):
  """The Pearson Csiszar-function in log-space.

  A Csiszar-function is a member of,

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

  The Pearson Csiszar-function is:

  ```none
  f(u) = (u - 1)**2
  ```

  Warning: this function makes non-log-space calculations and may therefore be
  numerically unstable for `|logu| >> 0`.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    Func_of_u: `float`-like `Tensor` of the Csiszar-function evaluated at
      `u = exp(logu)`.
  """

  with tf.compat.v1.name_scope(arg_1, "Func", [arg_0]):
    arg_0 = tf.convert_to_tensor(value=arg_0, arg_1="logu")
    return tf.square(tf.math.expm1(arg_0))


def Func(arg_0, arg_1=None):
  """The Squared-Hellinger Csiszar-function in log-space.

  A Csiszar-function is a member of,

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

  The Squared-Hellinger Csiszar-function is:

  ```none
  f(u) = (sqrt(u) - 1)**2
  ```

  This Csiszar-function induces a symmetric f-Divergence, i.e.,
  `D_f[p, q] = D_f[q, p]`.

  Warning: this function makes non-log-space calculations and may therefore be
  numerically unstable for `|logu| >> 0`.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    Func_of_u: `float`-like `Tensor` of the Csiszar-function
      evaluated at `u = exp(logu)`.
  """

  with tf.compat.v1.name_scope(arg_1, "Func", [arg_0]):
    arg_0 = tf.convert_to_tensor(value=arg_0, arg_1="logu")
    return pearson(0.5 * arg_0)


def Func(arg_0, arg_1, arg_2=False, arg_3=None):
  """The T-Power Csiszar-function in log-space.

  A Csiszar-function is a member of,

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

  When `self_normalized = True` the T-Power Csiszar-function is:

  ```none
  f(u) = s [ u**t - 1 - t(u - 1) ]
  s = { -1   0 < t < 1
      { +1   otherwise
  ```

  When `self_normalized = False` the `- t(u - 1)` term is omitted.

  This is similar to the `amari_alpha` Csiszar-function, with the associated
  divergence being the same up to factors depending only on `t`.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    t:  `Tensor` of same `dtype` as `logu` and broadcastable shape.
    self_normalized: Python `bool` indicating whether `f'(u=1)=0`.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    Func_of_u: `float`-like `Tensor` of the Csiszar-function evaluated
      at `u = exp(logu)`.
  """
  with tf.compat.v1.name_scope(arg_3, "Func", [arg_0, arg_1]):
    arg_0 = tf.convert_to_tensor(value=arg_0, arg_3="logu")
    arg_1 = tf.convert_to_tensor(value=arg_1, dtype=arg_0.dtype.base_dtype, arg_3="t")
    arg_4 = tf.math.expm1(arg_1 * arg_0)
    if arg_2:
      arg_4 -= arg_1 * tf.math.expm1(arg_0)
    arg_4 *= tf.where(tf.logical_and(0. < arg_1, arg_1 < 1.),
                   -tf.ones_like(arg_1),
                   tf.ones_like(arg_1))
    return arg_4


def Func(arg_0, arg_1=None):
  """The log1p-abs Csiszar-function in log-space.

  A Csiszar-function is a member of,

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

  The Log1p-Abs Csiszar-function is:

  ```none
  f(u) = u**(sign(u-1)) - 1
  ```

  This function is so-named because it was invented from the following recipe.
  Choose a convex function g such that g(0)=0 and solve for f:

  ```none
  log(1 + f(u)) = g(log(u)).
    <=>
  f(u) = exp(g(log(u))) - 1
  ```

  That is, the graph is identically `g` when y-axis is `log1p`-domain and x-axis
  is `log`-domain.

  Warning: this function makes non-log-space calculations and may therefore be
  numerically unstable for `|logu| >> 0`.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    Func_of_u: `float`-like `Tensor` of the Csiszar-function evaluated
      at `u = exp(logu)`.
  """

  with tf.compat.v1.name_scope(arg_1, "Func", [arg_0]):
    arg_0 = tf.convert_to_tensor(value=arg_0, arg_1="logu")
    return tf.math.expm1(tf.abs(arg_0))


def Func(arg_0, arg_1=None):
  """The Jeffreys Csiszar-function in log-space.

  A Csiszar-function is a member of,

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

  The Jeffreys Csiszar-function is:

  ```none
  f(u) = 0.5 ( u log(u) - log(u) )
       = 0.5 kl_forward + 0.5 kl_reverse
       = symmetrized_csiszar_function(kl_reverse)
       = symmetrized_csiszar_function(kl_forward)
  ```

  This Csiszar-function induces a symmetric f-Divergence, i.e.,
  `D_f[p, q] = D_f[q, p]`.

  Warning: this function makes non-log-space calculations and may therefore be
  numerically unstable for `|logu| >> 0`.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    Func_of_u: `float`-like `Tensor` of the Csiszar-function evaluated
      at `u = exp(logu)`.
  """

  with tf.compat.v1.name_scope(arg_1, "Func", [arg_0]):
    arg_0 = tf.convert_to_tensor(value=arg_0, arg_1="logu")
    return 0.5 * tf.math.expm1(arg_0) * arg_0


def Func(arg_0, arg_1=False, arg_2=None):
  """The Modified-GAN Csiszar-function in log-space.

  A Csiszar-function is a member of,

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

  When `self_normalized = True` the modified-GAN (Generative/Adversarial
  Network) Csiszar-function is:

  ```none
  f(u) = log(1 + u) - log(u) + 0.5 (u - 1)
  ```

  When `self_normalized = False` the `0.5 (u - 1)` is omitted.

  The unmodified GAN Csiszar-function is identical to Jensen-Shannon (with
  `self_normalized = False`).

  Warning: this function makes non-log-space calculations and may therefore be
  numerically unstable for `|logu| >> 0`.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    self_normalized: Python `bool` indicating whether `f'(u=1)=0`. When
      `f'(u=1)=0` the implied Csiszar f-Divergence remains non-negative even
      when `p, q` are unnormalized measures.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    chi_square_of_u: `float`-like `Tensor` of the Csiszar-function evaluated
      at `u = exp(logu)`.
  """

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


def Func(arg_0, arg_1, arg_2=None):
  """Calculates the dual Csiszar-function in log-space.

  A Csiszar-function is a member of,

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

  The Csiszar-dual is defined as:

  ```none
  f^*(u) = u f(1 / u)
  ```

  where `f` is some other Csiszar-function.

  For example, the dual of `kl_reverse` is `kl_forward`, i.e.,

  ```none
  f(u) = -log(u)
  f^*(u) = u f(1 / u) = -u log(1 / u) = u log(u)
  ```

  The dual of the dual is the original function:

  ```none
  f^**(u) = {u f(1/u)}^*(u) = u (1/u) f(1/(1/u)) = f(u)
  ```

  Warning: this function makes non-log-space calculations and may therefore be
  numerically unstable for `|logu| >> 0`.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    csiszar_function: Python `callable` representing a Csiszar-function over
      log-domain.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    dual_f_of_u: `float`-like `Tensor` of the result of calculating the dual of
      `f` at `u = exp(logu)`.
  """

  with tf.compat.v1.name_scope(arg_2, "Func", [arg_0]):
    return tf.exp(arg_0) * arg_1(-arg_0)


def Func(
    arg_0,
    arg_1,
    arg_2,
    arg_3,
    arg_4=None,
    arg_5=None,
    arg_6=None):
  """Monte-Carlo approximation of the Csiszar f-Divergence.

  A Csiszar-function is a member of,

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

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

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

  Tricks: Reparameterization and Score-Gradient

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

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

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

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

  Example Application:

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

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

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

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

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

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


def Func(arg_0, arg_1=None):
  """Helper to `csiszar_vimco`; computes `log_avg_u`, `log_sooavg_u`.

  `axis = 0` of `logu` is presumed to correspond to iid samples from `q`, i.e.,

  ```none
  logu[j] = log(u[j])
  u[j] = p(x, h[j]) / q(h[j] | x)
  h[j] iid~ q(H | x)
  ```

  Args:
    logu: Floating-type `Tensor` representing `log(p(x, h) / q(h | x))`.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    log_avg_u: `logu.dtype` `Tensor` corresponding to the natural-log of the
      average of `u`. The sum of the gradient of `log_avg_u` is `1`.
    log_sooavg_u: `logu.dtype` `Tensor` characterized by the natural-log of the
      average of `u`` except that the average swaps-out `u[i]` for the
      leave-`i`-out Geometric-average. The mean of the gradient of
      `log_sooavg_u` is `1`. Mathematically `log_sooavg_u` is,
      ```none
      log_sooavg_u[i] = log(Avg{h[j ; i] : j=0, ..., m-1})
      h[j ; i] = { u[j]                              j!=i
                 { GeometricAverage{u[k] : k != i}   j==i
      ```

  """
  with tf.compat.v1.name_scope(arg_1, "Func", [arg_0]):
    arg_0 = tf.convert_to_tensor(value=arg_0, arg_1="logu")

    arg_2 = tf.compat.dimension_value(arg_0.shape.with_rank_at_least(1)[0])
    if arg_2 is None:
      arg_2 = tf.shape(input=arg_0)[0]
      arg_3 = tf.math.log(tf.cast(arg_2, dtype=arg_0.dtype))
      arg_4 = tf.cast(arg_2 - 1, dtype=arg_0.dtype)
    else:
      arg_3 = np.log(arg_2).astype(arg_0.dtype.as_numpy_dtype)
      arg_4 = np.asarray(arg_2 - 1, dtype=arg_0.dtype.as_numpy_dtype)

    # Throughout we reduce across axis=0 since this is presumed to be iid
    # samples.

    arg_5 = tf.reduce_max(input_tensor=arg_0, axis=0)
    arg_6 = tf.reduce_logsumexp(
        input_tensor=arg_0 - arg_5, axis=0)

    # log_loosum_u[i] =
    # = logsumexp(logu[j] : j != i)
    # = log( exp(logsumexp(logu)) - exp(logu[i]) )
    # = log( exp(logsumexp(logu - logu[i])) exp(logu[i])  - exp(logu[i]))
    # = logu[i] + log(exp(logsumexp(logu - logu[i])) - 1)
    # = logu[i] + log(exp(logsumexp(logu) - logu[i]) - 1)
    # = logu[i] + softplus_inverse(logsumexp(logu) - logu[i])
    arg_7 = arg_6 + (arg_5 - arg_0)
    # We use `d != 0` rather than `d > 0.` because `d < 0.` should never
    # happens; if it does we want to complain loudly (which `softplus_inverse`
    # will).
    arg_8 = tf.not_equal(arg_7, 0.)
    arg_9 = tf.where(arg_8, arg_7, tf.ones_like(arg_7))
    arg_10 = arg_0 + tfd.softplus_inverse(arg_9)

    arg_11 = np.array(np.inf, dtype=arg_0.dtype.as_numpy_dtype)

    # When not(d_ok) and is_positive_and_largest then we manually compute the
    # log_loosum_u. (We can efficiently do this for any one point but not all,
    # hence we still need the above calculation.) This is good because when
    # this condition is met, we cannot use the above calculation; its -inf.
    # We now compute the log-leave-out-max-sum, replicate it to every
    # point and make sure to select it only when we need to.
    arg_12 = tf.logical_and(
        arg_0 > 0.,
        tf.equal(arg_0, arg_5[tf.newaxis, ...]))
    arg_13 = tf.reduce_logsumexp(
        input_tensor=tf.where(arg_12,
                              tf.fill(tf.shape(input=arg_0), -arg_11), arg_0),
        axis=0,
        keepdims=True)
    arg_13 = tf.tile(
        arg_13,
        multiples=1 + tf.pad(tensor=[arg_2 - 1], paddings=[[0, tf.rank(arg_0) - 1]]))

    arg_14 = tf.where(arg_12, arg_13,
                               tf.fill(tf.shape(input=arg_7), -arg_11))

    arg_15 = tf.where(arg_8, arg_10, arg_14)

    # The swap-one-out-sum ("soosum") is n different sums, each of which
    # replaces the i-th item with the i-th-left-out average, i.e.,
    # soo_sum_u[i] = [exp(logu) - exp(logu[i])] + exp(mean(logu[!=i]))
    #              =  exp(log_loosum_u[i])      + exp(looavg_logu[i])
    arg_16 = (tf.reduce_sum(input_tensor=arg_0, axis=0) - arg_0) / arg_4
    arg_17 = tf.reduce_logsumexp(
        input_tensor=tf.stack([arg_15, arg_16]), axis=0)

    arg_18 = arg_6 + arg_5 - arg_3
    arg_19 = arg_17 - arg_3

    arg_18.set_shape(arg_0.shape.with_rank_at_least(1)[1:])
    arg_19.set_shape(arg_0.shape)

    return arg_18, arg_19


def Func(arg_0, arg_1, arg_2):
  """Like batch_gather, but broadcasts to the left of axis."""
  # batch_gather assumes...
  #   params.shape =  [A1,...,AN, B1,...,BM]
  #   indices.shape = [A1,...,AN, C]
  # which gives output of shape
  #                   [A1,...,AN, C, B1,...,BM]
  # Here we broadcast dims of each to the left of `axis` in params, and left of
  # the rightmost dim in indices, e.g. we can
  # have
  #   params.shape =  [A1,...,AN, B1,...,BM]
  #   indices.shape = [a1,...,aN, C],
  # where ai broadcasts with Ai.

  # leading_bcast_shape is the broadcast of [A1,...,AN] and [a1,...,aN].
  arg_3 = tf.broadcast_dynamic_shape(
      tf.shape(input=arg_0)[:arg_2],
      tf.shape(input=arg_1)[:-1])
  arg_0 += tf.zeros(
      tf.concat((arg_3, tf.shape(input=arg_0)[arg_2:]), arg_2=0),
      dtype=arg_0.dtype)
  arg_1 += tf.zeros(
      tf.concat((arg_3, tf.shape(input=arg_1)[-1:]), arg_2=0),
      dtype=arg_1.dtype)
  return tf.compat.v1.batch_gather(arg_0, arg_1)


def Func(arg_0, arg_1, arg_2):
  """Broadcasts the event or distribution parameters."""
  if dtype_util.is_integer(arg_0.dtype):
    pass
  elif dtype_util.is_floating(arg_0.dtype):
    # When `validate_args=True` we've already ensured int/float casting
    # is closed.
    arg_0 = tf.cast(arg_0, dtype=tf.int32)
  else:
    raise TypeError("`value` should have integer `dtype` or "
                    "`self.dtype` ({})".format(arg_2))
  arg_3 = (
      tensorshape_util.rank(arg_1.shape) is not None and
      tensorshape_util.is_fully_defined(arg_1.shape[:-1]) and
      tensorshape_util.is_fully_defined(arg_0.shape))
  if not arg_3 or arg_1.shape[:-1] != arg_0.shape:
    arg_1 *= tf.ones_like(arg_0[..., tf.newaxis],
                           dtype=arg_1.dtype)
    arg_4 = tf.shape(input=arg_1)[:-1]
    arg_0 *= tf.ones(arg_4, dtype=arg_0.dtype)
    if tensorshape_util.rank(arg_1.shape) is not None:
      tensorshape_util.set_shape(arg_0, arg_1.shape[:-1])

  return arg_0, arg_1


def Func(
    arg_0,
    arg_1,
    arg_2,
    arg_3=None,
    arg_4=None,
    arg_5=None,
    arg_6='Func'):
  r"""Importance sampling with a positive function, in log-space.

  With \\(p(z) := exp^{log_p(z)}\\), and \\(f(z) = exp{log_f(z)}\\),
  this `Op` returns

  \\(Log[ n^{-1} sum_{i=1}^n [ f(z_i) p(z_i) / q(z_i) ] ],  z_i ~ q,\\)
  \\(\approx Log[ E_q[ f(Z) p(Z) / q(Z) ] ]\\)
  \\(=       Log[E_p[f(Z)]]\\)

  This integral is done in log-space with max-subtraction to better handle the
  often extreme values that `f(z) p(z) / q(z)` can take on.

  In contrast to `expectation_importance_sampler`, this `Op` returns values in
  log-space.

  User supplies either `Tensor` of samples `z`, or number of samples to draw `n`

  Args:
    log_f: Callable mapping samples from `sampling_dist_q` to `Tensors` with
      shape broadcastable to `q.batch_shape`.
      For example, `log_f` works "just like" `sampling_dist_q.log_prob`.
    log_p:  Callable mapping samples from `sampling_dist_q` to `Tensors` with
      shape broadcastable to `q.batch_shape`.
      For example, `log_p` works "just like" `q.log_prob`.
    sampling_dist_q:  The sampling distribution.
      `tfp.distributions.Distribution`.
      `float64` `dtype` recommended.
      `log_p` and `q` should be supported on the same set.
    z:  `Tensor` of samples from `q`, produced by `q.sample` for some `n`.
    n:  Integer `Tensor`.  Number of samples to generate if `z` is not provided.
    seed:  Python integer to seed the random number generator.
    name:  A name to give this `Op`.

  Returns:
    Logarithm of the importance sampling estimate.  `Tensor` with `shape` equal
      to batch shape of `q`, and `dtype` = `q.dtype`.
  """
  arg_7 = arg_2
  with tf.name_scope(arg_6):
    arg_3 = _get_samples(arg_7, arg_3, arg_4, arg_5)
    arg_8 = arg_0(arg_3) + arg_1(arg_3) - arg_7.log_prob(arg_3)
    return _logspace_mean(arg_8)


def Func(arg_0, arg_1, arg_2):
  """Broadcasts the event or samples."""
  # This is the shape of self.samples, without the samples axis, i.e. the shape
  # of the result of a call to dist.sample(). This way we can broadcast it with
  # event to get a properly-sized event, then add the singleton dim back at
  # -event_ndims - 1.
  arg_3 = tf.concat(
      [tf.shape(input=arg_1)[:-arg_2 - 1],
       tf.shape(input=arg_1)[tf.rank(arg_1) - arg_2:]],
      axis=0)
  arg_0 *= tf.ones(arg_3, dtype=arg_0.dtype)
  arg_0 = tf.expand_dims(arg_0, axis=-arg_2 - 1)
  arg_1 *= tf.ones_like(arg_0, dtype=arg_1.dtype)

  return arg_0, arg_1


def Func(arg_0,
             arg_1,
             arg_2=1e-8,
             arg_3=0,
             arg_4=0,
             arg_5=None,
             arg_6=50,
             arg_7=1,
             arg_8=None,
             arg_9=None):
  """Applies the BFGS algorithm to Func a differentiable function.

  Performs unconstrained minimization of a differentiable function using the
  BFGS scheme. For details of the algorithm, see [Nocedal and Wright(2006)][1].

  ### Usage:

  The following example demonstrates the BFGS optimizer attempting to find the
  minimum for a simple two dimensional quadratic objective function.

  ```python
    minimum = np.array([1.0, 1.0])  # The center of the quadratic bowl.
    scales = np.array([2.0, 3.0])  # The scales along the two axes.

    # The objective function and the gradient.
    def quadratic(x):
      value = tf.reduce_sum(scales * (x - minimum) ** 2)
      return value, tf.gradients(value, x)[0]

    start = tf.constant([0.6, 0.8])  # Starting point for the search.
    optim_results = tfp.optimizer.bfgs_Func(
        quadratic, initial_position=start, tolerance=1e-8)

    with tf.Session() as session:
      results = session.run(optim_results)
      # Check that the search converged
      assert(results.converged)
      # Check that the argmin is close to the actual value.
      np.testing.assert_allclose(results.position, minimum)
      # Print out the total number of function evaluations it took. Should be 6.
      print ("Function evaluations: %d" % results.num_objective_evaluations)
  ```

  ### References:
  [1]: Jorge Nocedal, Stephen Wright. Numerical Optimization. Springer Series in
    Operations Research. pp 136-140. 2006
    http://pages.mtu.edu/~struther/Courses/OLD/Sp2013/5630/Jorge_Nocedal_Numerical_optimization_267490.pdf

  Args:
    value_and_gradients_function:  A Python callable that accepts a point as a
      real `Tensor` and returns a tuple of `Tensor`s of real dtype containing
      the value of the function and its gradient at that point. The function
      to be Funcd. The input should be of shape `[..., n]`, where `n` is
      the size of the domain of input points, and all others are batching
      dimensions. The first component of the return value should be a real
      `Tensor` of matching shape `[...]`. The second component (the gradient)
      should also be of shape `[..., n]` like the input value to the function.
    initial_position: real `Tensor` of shape `[..., n]`. The starting point, or
      points when using batching dimensions, of the search procedure. At these
      points the function value and the gradient norm should be finite.
    tolerance: Scalar `Tensor` of real dtype. Specifies the gradient tolerance
      for the procedure. If the supremum norm of the gradient vector is below
      this number, the algorithm is stopped.
    x_tolerance: Scalar `Tensor` of real dtype. If the absolute change in the
      position between one iteration and the next is smaller than this number,
      the algorithm is stopped.
    f_relative_tolerance: Scalar `Tensor` of real dtype. If the relative change
      in the objective value between one iteration and the next is smaller
      than this value, the algorithm is stopped.
    initial_inverse_hessian_estimate: Optional `Tensor` of the same dtype
      as the components of the output of the `value_and_gradients_function`.
      If specified, the shape should broadcastable to shape `[..., n, n]`; e.g.
      if a single `[n, n]` matrix is provided, it will be automatically
      broadcasted to all batches. Alternatively, one can also specify a
      different hessian estimate for each batch member.
      For the correctness of the algorithm, it is required that this parameter
      be symmetric and positive definite. Specifies the starting estimate for
      the inverse of the Hessian at the initial point. If not specified,
      the identity matrix is used as the starting estimate for the
      inverse Hessian.
    max_iterations: Scalar positive int32 `Tensor`. The maximum number of
      iterations for BFGS updates.
    parallel_iterations: Positive integer. The number of iterations allowed to
      run in parallel.
    stopping_condition: (Optional) A Python function that takes as input two
      Boolean tensors of shape `[...]`, and returns a Boolean scalar tensor.
      The input tensors are `converged` and `failed`, indicating the current
      status of each respective batch member; the return value states whether
      the algorithm should stop. The default is tfp.optimizer.converged_all
      which only stops when all batch members have either converged or failed.
      An alternative is tfp.optimizer.converged_any which stops as soon as one
      batch member has converged, or when all have failed.
    name: (Optional) Python str. The name prefixed to the ops created by this
      function. If not supplied, the default name 'Func' is used.

  Returns:
    optimizer_results: A namedtuple containing the following items:
      converged: boolean tensor of shape `[...]` indicating for each batch
        member whether the minimum was found within tolerance.
      failed:  boolean tensor of shape `[...]` indicating for each batch
        member whether a line search step failed to find a suitable step size
        satisfying Wolfe conditions. In the absence of any constraints on the
        number of objective evaluations permitted, this value will
        be the complement of `converged`. However, if there is
        a constraint and the search stopped due to available
        evaluations being exhausted, both `failed` and `converged`
        will be simultaneously False.
      num_objective_evaluations: The total number of objective
        evaluations performed.
      position: A tensor of shape `[..., n]` containing the last argument value
        found during the search from each starting point. If the search
        converged, then this value is the argmin of the objective function.
      objective_value: A tensor of shape `[...]` with the value of the
        objective function at the `position`. If the search converged, then
        this is the (local) minimum of the objective function.
      objective_gradient: A tensor of shape `[..., n]` containing the gradient
        of the objective function at the `position`. If the search converged
        the max-norm of this tensor should be below the tolerance.
      inverse_hessian_estimate: A tensor of shape `[..., n, n]` containing the
        inverse of the estimated Hessian.
  """
  with tf.compat.v1.name_scope(
      arg_9, 'Func',
      [arg_1, arg_2, arg_5]):
    arg_1 = tf.convert_to_tensor(
        value=arg_1, arg_9='initial_position')
    arg_10 = arg_1.dtype.base_dtype
    arg_2 = tf.convert_to_tensor(
        value=arg_2, arg_10=arg_10, arg_9='grad_tolerance')
    arg_4 = tf.convert_to_tensor(
        value=arg_4, arg_10=arg_10, arg_9='f_relative_tolerance')
    arg_3 = tf.convert_to_tensor(
        value=arg_3, arg_10=arg_10, arg_9='x_tolerance')
    arg_6 = tf.convert_to_tensor(
        value=arg_6, arg_9='max_iterations')

    arg_11 = distribution_util.prefer_static_shape(arg_1)
    arg_12, arg_13 = arg_11[:-1], arg_11[-1]

    if arg_8 is None:
      arg_8 = bfgs_utils.converged_all

    # Control inputs are an optional list of tensors to evaluate before
    # the start of the search procedure. These can be used to assert the
    # validity of inputs to the search procedure.
    arg_14 = None

    if arg_5 is None:
      # Create a default initial inverse Hessian.
      arg_15 = tf.eye(arg_13,
                                   arg_12=arg_12,
                                   arg_10=arg_10,
                                   arg_9='initial_inv_hessian')
    else:
      # If an initial inverse Hessian is supplied, compute some control inputs
      # to ensure that it is positive definite and symmetric.
      arg_15 = tf.convert_to_tensor(
          value=arg_5,
          arg_10=arg_10,
          arg_9='initial_inv_hessian')
      arg_14 = _inv_hessian_control_inputs(arg_15)
      arg_16 = tf.concat([arg_12, [arg_13, arg_13]], 0)
      arg_15 = tf.broadcast_to(arg_15, arg_16)

    # The `state` here is a `BfgsOptimizerResults` tuple with values for the
    # current state of the algorithm computation.
    def _cond(arg_17):
      """Continue if iterations remain and stopping condition is not met."""
      return ((arg_17.num_iterations < arg_6) &
              tf.logical_not(arg_8(arg_17.converged, arg_17.failed)))

    def _body(arg_17):
      """Main optimization loop."""
      arg_18 = _get_search_direction(arg_17.inverse_hessian_estimate,
                                               arg_17.objective_gradient)
      arg_19 = tf.reduce_sum(
          input_tensor=arg_17.objective_gradient * arg_18, axis=-1)

      # If the derivative at the start point is not negative, recompute the
      # search direction with the initial inverse Hessian.
      arg_20 = (~arg_17.failed & ~arg_17.converged &
                     (arg_19 >= 0))

      arg_21 = _get_search_direction(
          arg_15, arg_17.objective_gradient)

      arg_22 = tf.where(
          arg_20, arg_21, arg_18)
      arg_23 = tf.where(
          arg_20, arg_15, arg_17.inverse_hessian_estimate)

      # Replace the hessian estimate in the state, in case it had to be reset.
      arg_24 = bfgs_utils.update_fields(
          arg_17, inverse_hessian_estimate=arg_23)

      arg_25 = bfgs_utils.line_search_step(
          arg_24,
          arg_0, arg_22,
          arg_2, arg_4, arg_3, arg_8)

      # Update the inverse Hessian if needed and continue.
      return [_update_inv_hessian(arg_24, arg_25)]

    arg_26 = bfgs_utils.get_initial_state_args(
        arg_0,
        arg_1,
        arg_2,
        arg_14)
    arg_26['inverse_hessian_estimate'] = arg_15
    arg_27 = BfgsOptimizerResults(**arg_26)
    return tf.while_loop(
        cond=_cond,
        body=_body,
        loop_vars=[arg_27],
        arg_7=arg_7)[0]


def Func(arg_0):
  """Computes control inputs to validate a provided inverse Hessian.

  These ensure that the provided inverse Hessian is positive definite and
  symmetric.

  Args:
    inv_hessian: The starting estimate for the inverse of the Hessian at the
      initial point.

  Returns:
    A list of tf.Assert ops suitable for use with tf.control_dependencies.
  """
  # The easiest way to validate if the inverse Hessian is positive definite is
  # to compute its Cholesky decomposition.
  arg_1 = tf.reduce_all(
      input_tensor=tf.math.is_finite(tf.linalg.cholesky(arg_0)),
      axis=[-1, -2])

  # Then check that the supplied inverse Hessian is symmetric.
  arg_2 = tf.equal(bfgs_utils.norm(
      arg_0 - _batch_transpose(arg_0), dims=2), 0)

  # Simply adding a control dependencies on these results is not enough to
  # trigger them, we need to add asserts on the results.
  return [tf.Assert(arg_1,
                    ['Initial inverse Hessian is not positive definite.',
                     arg_0]),
          tf.Assert(arg_2,
                    ['Initial inverse Hessian is not symmetric',
                     arg_0])]


def Func(arg_0, arg_1):
  """Update the BGFS state by computing the next inverse hessian estimate."""
  # Only update the inverse Hessian if not already failed or converged.
  arg_2 = ~arg_1.converged & ~arg_1.failed

  # Compute the normalization term (y^T . s), should not update if is singular.
  arg_3 = arg_1.objective_gradient - arg_0.objective_gradient
  arg_4 = arg_1.position - arg_0.position
  arg_5 = tf.reduce_sum(
      input_tensor=arg_3 * arg_4, axis=-1)
  arg_2 = arg_2 & ~tf.equal(arg_5, 0)

  def _doFunc():
    arg_6 = _bfgs_inv_hessian_update(
        arg_3, arg_4, arg_5,
        arg_0.inverse_hessian_estimate)
    return bfgs_utils.update_fields(
        arg_1,
        inverse_hessian_estimate=tf.where(arg_2,
                                          arg_6,
                                          arg_0.inverse_hessian_estimate))

  return prefer_static.cond(
      tf.reduce_any(input_tensor=arg_2),
      _doFunc,
      lambda: arg_1)


def Func(arg_0, arg_1, arg_2,
                             arg_3):
  """Applies the BFGS update to the inverse Hessian estimate.

  The BFGS update rule is (note A^T denotes the transpose of a vector/matrix A).

  ```None
    rho = 1/(grad_delta^T * position_delta)
    U = (I - rho * position_delta * grad_delta^T)
    H_1 =  U * H_0 * U^T + rho * position_delta * position_delta^T
  ```

  Here, `H_0` is the inverse Hessian estimate at the previous iteration and
  `H_1` is the next estimate. Note that `*` should be interpreted as the
  matrix multiplication (with the understanding that matrix multiplication for
  scalars is usual multiplication and for matrix with vector is the action of
  the matrix on the vector.).

  The implementation below utilizes an expanded version of the above formula
  to avoid the matrix multiplications that would be needed otherwise. By
  expansion it is easy to see that one only needs matrix-vector or
  vector-vector operations. The expanded version is:

  ```None
    f = 1 + rho * (grad_delta^T * H_0 * grad_delta)
    H_1 - H_0 = - rho * [position_delta * (H_0 * grad_delta)^T +
                        (H_0 * grad_delta) * position_delta^T] +
                  rho * f * [position_delta * position_delta^T]
  ```

  All the terms in square brackets are matrices and are constructed using
  vector outer products. All the other terms on the right hand side are scalars.
  Also worth noting that the first and second lines are both rank 1 updates
  applied to the current inverse Hessian estimate.

  Args:
    grad_delta: Real `Tensor` of shape `[..., n]`. The difference between the
      gradient at the new position and the old position.
    position_delta: Real `Tensor` of shape `[..., n]`. The change in position
      from the previous iteration to the current one.
    normalization_factor: Real `Tensor` of shape `[...]`. Should be equal to
      `grad_delta^T * position_delta`, i.e. `1/rho` as defined above.
    inv_hessian_estimate: Real `Tensor` of shape `[..., n, n]`. The previous
      estimate of the inverse Hessian. Should be positive definite and
      symmetric.

  Returns:
    A tuple containing the following fields
      is_valid: A Boolean `Tensor` of shape `[...]` indicating batch members
        where the update succeeded. The update can fail if the position change
        becomes orthogonal to the gradient change.
      next_inv_hessian_estimate: A `Tensor` of shape `[..., n, n]`. The next
        Hessian estimate updated using the BFGS update scheme. If the
        `inv_hessian_estimate` is symmetric and positive definite, the
        `next_inv_hessian_estimate` is guaranteed to satisfy the same
        conditions.
  """
  # The quadratic form: y^T.H.y; where H is the inverse Hessian and y is the
  # gradient change.
  arg_4 = _mul_right(arg_3, arg_0)
  arg_5 = tf.reduce_sum(
      input_tensor=arg_4 * arg_0, axis=-1)

  # The first rank 1 update term requires the outer product: s.y^T.
  arg_6 = _tensor_product(arg_1, arg_4)

  def _expand_scalar(arg_7):
    # Expand dimensions of a batch of scalars to multiply or divide a matrix.
    return arg_7[..., tf.newaxis, tf.newaxis]

  # Symmetrize
  arg_6 += _tensor_product(arg_4, arg_1)
  arg_8 = _tensor_product(arg_1, arg_1)
  with tf.control_dependencies([arg_8]):
    arg_8 *= _expand_scalar(
        1 + arg_5 / arg_2)

  return (arg_3 +
          (arg_8 - arg_6) / _expand_scalar(arg_2))


def Func(arg_0):
  """Transpose a possibly batched matrix.

  Args:
    mat: A `tf.Tensor` of shape `[..., n, m]`.

  Returns:
    A tensor of shape `[..., m, n]` with matching batch dimensions.
  """
  arg_1 = distribution_util.prefer_static_rank(arg_0)
  arg_2 = tf.range(arg_1)
  arg_2 = tf.concat([arg_2[:-2], [arg_2[-1], arg_2[-2]]], axis=0)
  return tf.transpose(a=arg_0, arg_2=arg_2)


def Func(arg_0, arg_1):
  """Maybe add `ndims` ones to `x.shape` on the right.

  If `ndims` is zero, this is a no-op; otherwise, we will create and return a
  new `Tensor` whose shape is that of `x` with `ndims` ones concatenated on the
  right side. If the shape of `x` is known statically, the shape of the return
  value will be as well.

  Args:
    x: The `Tensor` we'll return a reshaping of.
    ndims: Python `integer` number of ones to pad onto `x.shape`.
  Returns:
    If `ndims` is zero, `x`; otherwise, a `Tensor` whose shape is that of `x`
    with `ndims` ones concatenated on the right side. If possible, returns a
    `Tensor` whose shape is known statically.
  Raises:
    ValueError: if `ndims` is not a Python `integer` greater than or equal to
    zero.
  """
  if not (isinstance(arg_1, int) and arg_1 >= 0):
    raise ValueError(
        '`ndims` must be a Python `integer` greater than zero. Got: {}'
        .format(arg_1))
  if arg_1 == 0:
    return arg_0
  arg_0 = tf.convert_to_tensor(value=arg_0)
  arg_2 = arg_0.shape
  arg_3 = distribution_util.pad(
      tf.shape(input=arg_0), axis=0, back=True, value=1, count=arg_1)
  arg_0 = tf.reshape(arg_0, arg_3)
  arg_0.set_shape(arg_2.concatenate([1]*arg_1))
  return arg_0


def Func(arg_0, arg_1):
  """Return `Tensor` with right-most ndims summed.

  Args:
    x: the `Tensor` whose right-most `ndims` dimensions to sum
    ndims: number of right-most dimensions to sum.

  Returns:
    A `Tensor` resulting from calling `reduce_sum` on the `ndims` right-most
    dimensions. If the shape of `x` is statically known, the result will also
    have statically known shape. Otherwise, the resulting shape will only be
    known at runtime.
  """
  arg_0 = tf.convert_to_tensor(value=arg_0)
  if arg_0.shape.ndims is not None:
    arg_2 = tf.range(arg_0.shape.ndims - arg_1, arg_0.shape.ndims)
  else:
    arg_2 = tf.range(tf.rank(arg_0) - arg_1, tf.rank(arg_0))
  return tf.reduce_sum(input_tensor=arg_0, axis=arg_2)


def Func(arg_0, arg_1=None):
  """A sqrt function whose gradient at zero is very large but finite.

  Args:
    x: a `Tensor` whose sqrt is to be computed.
    name: a Python `str` prefixed to all ops created by this function.
      Default `None` (i.e., "Func").

  Returns:
    sqrt: the square root of `x`, with an overridden gradient at zero
    grad: a gradient function, which is the same as sqrt's gradient everywhere
      except at zero, where it is given a large finite value, instead of `inf`.

  Raises:
    TypeError: if `tf.convert_to_tensor(x)` is not a `float` type.

  Often in kernel functions, we need to compute the L2 norm of the difference
  between two vectors, `x` and `y`: `sqrt(sum_i((x_i - y_i) ** 2))`. In the
  case where `x` and `y` are identical, e.g., on the diagonal of a kernel
  matrix, we get `NaN`s when we take gradients with respect to the inputs. To
  see, this consider the forward pass:

    ```
    [x_1 ... x_N]  -->  [x_1 ** 2 ... x_N ** 2]  -->
        (x_1 ** 2 + ... + x_N ** 2)  -->  sqrt((x_1 ** 2 + ... + x_N ** 2))
    ```

  When we backprop through this forward pass, the `sqrt` yields an `inf` because
  `grad_z(sqrt(z)) = 1 / (2 * sqrt(z))`. Continuing the backprop to the left, at
  the `x ** 2` term, we pick up a `2 * x`, and when `x` is zero, we get
  `0 * inf`, which is `NaN`.

  We'd like to avoid these `NaN`s, since they infect the rest of the connected
  computation graph. Practically, when two inputs to a kernel function are
  equal, we are in one of two scenarios:
    1. We are actually computing k(x, x), in which case norm(x - x) is
       identically zero, independent of x. In this case, we'd like the
       gradient to reflect this independence: it should be zero.
    2. We are computing k(x, y), and x just *happens* to have the same value
       as y. The gradient at such inputs is in fact ill-defined (there is a
       cusp in the sqrt((x - y) ** 2) surface along the line x = y). There are,
       however, an infinite number of sub-gradients, all of which are valid at
       all such inputs. By symmetry, there is exactly one which is "special":
       zero, and we elect to use that value here. In practice, having two
       identical inputs to a kernel matrix is probably a pathological
       situation to be avoided, but that is better resolved at a higher level
       than this.

  To avoid the infinite gradient at zero, we use tf.custom_gradient to redefine
  the gradient at zero. We assign it to be a very large value, specifically
  the sqrt of the max value of the floating point dtype of the input. We use
  the sqrt (as opposed to just using the max floating point value) to avoid
  potential overflow when combining this value with others downstream.
  """
  with tf.compat.v1.name_scope(arg_1, 'Func', [arg_0]):
    arg_0 = tf.convert_to_tensor(value=arg_0, arg_1='x')
    if not arg_0.dtype.is_floating:
      raise TypeError('Input `x` must be floating type.')
    def grad(arg_2):
      arg_3 = np.sqrt(np.finfo(arg_0.dtype.as_numpy_dtype()).max)
      arg_4 = tf.where(
          tf.equal(arg_0, 0), tf.fill(tf.shape(input=arg_0), arg_3),
          0.5 * tf.math.rsqrt(arg_0))
      return arg_2 * arg_4
    return tf.sqrt(arg_0), grad