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

  Performs unconstrained minimization of a differentiable function using the
  L-BFGS scheme. See [Nocedal and Wright(2006)][1] for details of the algorithm.

  ### Usage:

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

  ```python
    # A high-dimensional quadratic bowl.
    ndims = 60
    minimum = np.ones([ndims], dtype='float64')
    scales = np.arange(ndims, dtype='float64') + 1.0

    # 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 = np.arange(ndims, 0, -1, dtype='float64')
    optim_results = tfp.optimizer.lbfgs_Func(
        quadratic, initial_position=start, num_correction_pairs=10,
        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)
  ```

  ### References:

  [1] Jorge Nocedal, Stephen Wright. Numerical Optimization. Springer Series
      in Operations Research. pp 176-180. 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 is 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 is a real `Tensor` of matching
      shape `[...]`. The second component (the gradient) is also 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.
    num_correction_pairs: Positive integer. Specifies the maximum number of
      (position_delta, gradient_delta) correction pairs to keep as implicit
      approximation of the Hessian matrix.
    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: None. Option currently not supported.
    max_iterations: Scalar positive int32 `Tensor`. The maximum number of
      iterations for L-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: Scalar boolean tensor indicating whether the minimum was
        found within tolerance.
      failed:  Scalar boolean tensor indicating 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 containing the last argument value found
        during the search. If the search converged, then
        this value is the argmin of the objective function.
      objective_value: A tensor containing 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 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.
      position_deltas: A tensor encoding information about the latest
        changes in `position` during the algorithm execution.
      gradient_deltas: A tensor encoding information about the latest
        changes in `objective_gradient` during the algorithm execution.
  """
  if arg_6 is not None:
    raise NotImplementedError(
        'Support of initial_inverse_hessian_estimate arg not yet implemented')

  if arg_9 is None:
    arg_9 = bfgs_utils.converged_all

  with tf.compat.v1.name_scope(arg_10, 'Func', [arg_1, arg_3]):
    arg_1 = tf.convert_to_tensor(
        value=arg_1, arg_10='initial_position')
    arg_11 = arg_1.dtype.base_dtype
    arg_3 = tf.convert_to_tensor(
        value=arg_3, arg_11=arg_11, arg_10='grad_tolerance')
    arg_5 = tf.convert_to_tensor(
        value=arg_5, arg_11=arg_11, arg_10='f_relative_tolerance')
    arg_4 = tf.convert_to_tensor(
        value=arg_4, arg_11=arg_11, arg_10='x_tolerance')
    arg_7 = tf.convert_to_tensor(
        value=arg_7, arg_10='max_iterations')

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

    def _body(arg_13):
      """Main optimization loop."""
      arg_14 = _get_search_direction(arg_13)

      # TODO(b/120134934): Check if the derivative at the start point is not
      # negative, if so then reset position/gradient deltas and recompute
      # search direction.

      arg_15 = bfgs_utils.line_search_step(
          arg_13,
          arg_0, arg_14,
          arg_3, arg_5, arg_4, arg_9)

      # If not failed or converged, update the Hessian estimate.
      arg_16 = ~(arg_15.converged | arg_15.failed)
      arg_17 = bfgs_utils.update_fields(
          arg_15,
          position_deltas=_queue_push(
              arg_13.position_deltas, arg_16,
              arg_15.position - arg_13.position),
          gradient_deltas=_queue_push(
              arg_13.gradient_deltas, arg_16,
              arg_15.objective_gradient - arg_13.objective_gradient))
      return [arg_17]

    arg_18 = _get_initial_state(arg_0,
                                       arg_1,
                                       arg_2,
                                       arg_3)
    return tf.while_loop(
        cond=_cond,
        body=_body,
        loop_vars=[arg_18],
        arg_8=arg_8)[0]


def Func(arg_0,
                       arg_1,
                       arg_2,
                       arg_3):
  """Create LBfgsOptimizerResults with initial state of search procedure."""
  arg_4 = bfgs_utils.get_initial_state_args(
      arg_0,
      arg_1,
      arg_3)
  arg_5 = _make_empty_queue_for(arg_2, arg_1)
  arg_4.update(position_deltas=arg_5, gradient_deltas=arg_5)
  return LBfgsOptimizerResults(**arg_4)


def Func(arg_0):
  """Computes the search direction to follow at the current state.

  On the `k`-th iteration of the main L-BFGS algorithm, the state has collected
  the most recent `m` correction pairs in position_deltas and gradient_deltas,
  where `k = state.num_iterations` and `m = min(k, num_correction_pairs)`.

  Assuming these, the code below is an implementation of the L-BFGS two-loop
  recursion algorithm given by [Nocedal and Wright(2006)][1]:

  ```None
    q_direction = objective_gradient
    for i in reversed(range(m)):  # First loop.
      inv_rho[i] = gradient_deltas[i]^T * position_deltas[i]
      alpha[i] = position_deltas[i]^T * q_direction / inv_rho[i]
      q_direction = q_direction - alpha[i] * gradient_deltas[i]

    kth_inv_hessian_factor = (gradient_deltas[-1]^T * position_deltas[-1] /
                              gradient_deltas[-1]^T * gradient_deltas[-1])
    r_direction = kth_inv_hessian_factor * I * q_direction

    for i in range(m):  # Second loop.
      beta = gradient_deltas[i]^T * r_direction / inv_rho[i]
      r_direction = r_direction + position_deltas[i] * (alpha[i] - beta)

    return -r_direction  # Approximates - H_k * objective_gradient.
  ```

  Args:
    state: A `LBfgsOptimizerResults` tuple with the current state of the
      search procedure.

  Returns:
    A real `Tensor` of the same shape as the `state.position`. The direction
    along which to perform line search.
  """
  # The number of correction pairs that have been collected so far.
  arg_1 = tf.minimum(
      arg_0.num_iterations,
      distribution_util.prefer_static_shape(arg_0.position_deltas)[0])

  def _two_loop_algorithm():
    """L-BFGS two-loop algorithm."""
    # Correction pairs are always appended to the end, so only the latest
    # `num_elements` vectors have valid position/gradient deltas.
    arg_2 = arg_0.position_deltas[-arg_1:]
    arg_3 = arg_0.gradient_deltas[-arg_1:]

    # Pre-compute all `inv_rho[i]`s.
    arg_4 = tf.reduce_sum(
        input_tensor=arg_3 * arg_2, axis=-1)

    def first_loop(arg_5, arg_6):
      arg_7, arg_8 = arg_5
      arg_9, arg_10, arg_11 = arg_6
      arg_12 = tf.reduce_sum(
          input_tensor=arg_9 * arg_8, axis=-1) / arg_11
      arg_13 = tf.expand_dims(arg_12, axis=-1) * arg_10
      return (arg_12, arg_8 - arg_13)

    # Run first loop body computing and collecting `alpha[i]`s, while also
    # computing the updated `q_direction` at each step.
    arg_14 = tf.zeros_like(arg_4[0])
    arg_15, arg_16 = tf.scan(
        first_loop, [arg_2, arg_3, arg_4],
        initializer=(arg_14, arg_0.objective_gradient), reverse=True)

    # We use `H^0_k = gamma_k * I` as an estimate for the initial inverse
    # hessian for the k-th iteration; then `r_direction = H^0_k * q_direction`.
    arg_17 = arg_4[-1] / tf.reduce_sum(
        input_tensor=arg_3[-1] * arg_3[-1], axis=-1)
    arg_18 = tf.expand_dims(arg_17, axis=-1) * arg_16[0]

    def second_loop(arg_18, arg_6):
      arg_12, arg_9, arg_10, arg_11 = arg_6
      arg_19 = tf.reduce_sum(
          input_tensor=arg_10 * arg_18, axis=-1) / arg_11
      arg_13 = tf.expand_dims(arg_12 - arg_19, axis=-1) * arg_9
      return arg_18 + arg_13

    # Finally, run second loop body computing the updated `r_direction` at each
    # step.
    arg_20 = tf.scan(
        second_loop, [arg_15, arg_2, arg_3, arg_4],
        initializer=arg_18)
    return -arg_20[-1]

  return prefer_static.cond(tf.equal(arg_1, 0),
                            (lambda: -arg_0.objective_gradient),
                            _two_loop_algorithm)


def Func(arg_0, arg_1):
  """Creates a `tf.Tensor` suitable to hold `k` element-shaped tensors.

  For example:

  ```python
    element = tf.constant([[0., 1., 2., 3., 4.],
                           [5., 6., 7., 8., 9.]])

    # A queue capable of holding 3 elements.
    Func(3, element)
    # => [[[ 0.,  0.,  0.,  0.,  0.],
    #      [ 0.,  0.,  0.,  0.,  0.]],
    #
    #     [[ 0.,  0.,  0.,  0.,  0.],
    #      [ 0.,  0.,  0.,  0.,  0.]],
    #
    #     [[ 0.,  0.,  0.,  0.,  0.],
    #      [ 0.,  0.,  0.,  0.,  0.]]]
  ```

  Args:
    k: A positive scalar integer, number of elements that each queue will hold.
    element: A `tf.Tensor`, only its shape and dtype information are relevant.

  Returns:
    A zero-filed `tf.Tensor` of shape `(k,) + tf.shape(element)` and same dtype
    as `element`.
  """
  arg_2 = tf.concat(
      [[arg_0], distribution_util.prefer_static_shape(arg_1)], axis=0)
  return tf.zeros(arg_2, dtype=arg_1.dtype.base_dtype)


def Func(arg_0):
  """Computes whether each square matrix in the input is positive semi-definite.

  Args:
    x: A floating-point `Tensor` of shape `[B1, ..., Bn, M, M]`.

  Returns:
    mask: A floating-point `Tensor` of shape `[B1, ... Bn]`.  Each
      scalar is 1 if the corresponding matrix was PSD, otherwise 0.
  """
  # Allegedly
  # https://scicomp.stackexchange.com/questions/12979/testing-if-a-matrix-is-positive-semi-definite
  # it is more efficient to test for positive semi-definiteness by
  # trying to compute the Cholesky decomposition -- the matrix is PSD
  # if you succeed and not PSD if you fail.  However, TensorFlow's
  # Cholesky raises an exception if _any_ of the input matrices are
  # not PSD, from which I don't know how to extract _which ones_, so I
  # proceed by explicitly computing all the eigenvalues and checking
  # whether they are all positive or not.
  #
  # Also, as was discussed in the answer, it is somewhat dangerous to
  # treat SPD-ness as binary in floating-point arithmetic. Cholesky
  # factorization can complete and 'look' like everything is fine
  # (e.g., O(1) entries and a diagonal of all ones) but the matrix can
  # have an exponential condition number.
  arg_1, arg_2 = tf.linalg.eigh(arg_0)
  return tf.cast(
      tf.reduce_min(input_tensor=arg_1, axis=-1) >= 0, dtype=arg_0.dtype)


def Func(
    arg_0, arg_1, arg_2, arg_3, arg_4):
  """Returns rejection samples from trying to get good correlation matrices.

  The proposal being rejected from is the uniform distribution on
  "correlation-like" matrices.  We say a matrix is "correlation-like"
  if it is a symmetric square matrix with all entries between -1 and 1
  (inclusive) and 1s on the main diagonal.  Of these, the ones that
  are positive semi-definite are exactly the correlation matrices.

  The rejection algorithm, then, is to sample a `Tensor` of
  `sample_shape` correlation-like matrices of dimensions `dim` by
  `dim`, and check each one for (i) being a correlation matrix (i.e.,
  PSD), and (ii) having determinant at least the corresponding entry
  of `det_bounds`.

  Args:
    det_bounds: A `Tensor` of lower bounds on the determinants of
      acceptable matrices.  The shape must broadcast with `sample_shape`.
    dim: A Python `int` dimension of correlation matrices to sample.
    sample_shape: Python `tuple` of `int` shape of the samples to
      compute, excluding the two matrix dimensions.
    dtype: The `dtype` in which to do the computation.
    seed: Random seed.

  Returns:
    weights: A `Tensor` of shape `sample_shape`.  Each entry is 0 if the
      corresponding matrix was not a correlation matrix, or had too
      small of a determinant.  Otherwise, the entry is the
      multiplicative inverse of the density of proposing that matrix
      uniformly, i.e., the volume of the set of `dim` by `dim`
      correlation-like matrices.
    volume: The volume of the set of `dim` by `dim` correlation-like
      matrices.
  """
  with tf.compat.v1.name_scope("rejection_sampler"):
    arg_5 = _uniform_correlation_like_matrix(
        arg_1, arg_2, arg_3, arg_4=arg_4)
    arg_6 = 2. ** (arg_1 * (arg_1 - 1) / 2.)
    # The density of proposing any given point is 1 / rej_proposal_volume;
    # The weight of that point should be scaled by
    # 1 / density = rej_proposal_volume.
    arg_7 = arg_6 * _psd_mask(
        arg_5) * _det_large_enough_mask(arg_5, arg_0)
    return arg_7, arg_6


def Func(arg_0, arg_1):
  """Computes a confidence interval for the mean of the given 1-D distribution.

  Assumes (and checks) that the given distribution is Bernoulli, i.e.,
  takes only two values.  This licenses using the CDF of the binomial
  distribution for the confidence, which is tighter (for extreme
  probabilities) than the DKWM inequality.  The method is known as the
  [Clopper-Pearson method]
  (https://en.wikipedia.org/wiki/Binomial_proportion_confidence_interval).

  Assumes:

  - The given samples were drawn iid from the distribution of interest.

  - The given distribution is a Bernoulli, i.e., supported only on
    low and high.

  Guarantees:

  - The probability (over the randomness of drawing the given sample)
    that the true mean is outside the returned interval is no more
    than the given error_rate.

  Args:
    samples: `np.ndarray` of samples drawn iid from the distribution
      of interest.
    error_rate: Python `float` admissible rate of mistakes.

  Returns:
    low: Lower bound of confidence interval.
    high: Upper bound of confidence interval.

  Raises:
    ValueError: If `samples` has rank other than 1 (batch semantics
      are not implemented), or if `samples` contains values other than
      `low` or `high` (as that makes the distribution not Bernoulli).
  """
  # TODO(b/78025336) Migrate this confidence interval function
  # to statistical_testing.py.  In order to do that
  # - Get the binomial CDF from the Binomial distribution
  # - Implement scalar root finding in TF.  Batch bisection search
  #   shouldn't be too hard, and is definitely good enough for this
  #   problem.  Batching the Brent algorithm (from scipy) that is used
  #   here may be more involved, but may also not be necessary---it's
  #   only used here because scipy made it convenient.  In particular,
  #   robustness is more important than speed here, which may make
  #   bisection search actively better.
  # - The rest is just a matter of rewriting in the appropriate style.
  if optimize is None or stats is None:
    raise ValueError(
        "Scipy is required for computing Clopper-Pearson confidence intervals")
  if len(arg_0.shape) != 1:
    raise ValueError("Batch semantics not implemented")
  arg_2 = len(arg_0)
  arg_3 = np.amin(arg_0)
  arg_4 = np.amax(arg_0)
  arg_5 = np.count_nonzero(arg_0 - arg_3)
  arg_6 = np.count_nonzero(arg_0 - arg_4)
  if arg_5 + arg_6 != arg_2:
    arg_7 = np.unique(arg_0)
    arg_8 = ("Purportedly Bernoulli distribution had distinct samples"
           " {}, {}, and {}".format(arg_7[0], arg_7[1], arg_7[2]))
    raise ValueError(arg_8)
  def p_small_enough(arg_9):
    arg_10 = stats.binom.logcdf(arg_5, arg_2, arg_9)
    return arg_10 - np.log(arg_1 / 2.)
  def p_big_enough(arg_9):
    arg_10 = stats.binom.logsf(arg_5, arg_2, arg_9)
    return arg_10 - np.log(arg_1 / 2.)
  arg_11 = optimize.brentq(
      p_small_enough, float(arg_5) / arg_2, 1., rtol=1e-9)
  arg_12 = optimize.brentq(
      p_big_enough, 0., float(arg_5) / arg_2, rtol=1e-9)
  arg_13 = arg_3 + (arg_4 - arg_3) * arg_12
  arg_14 = arg_3 + (arg_4 - arg_3) * arg_11
  return (arg_13, arg_14)


def Func(arg_0, arg_1, arg_2, arg_3):
  """Computes the von Mises CDF and its derivative via series expansion."""
  # Keep the number of terms as a float. It should be a small integer, so
  # exactly representable as a float.
  arg_2 = tf.cast(arg_2, arg_3=arg_3)

  def loop_body(arg_4, arg_5, arg_6, arg_7, arg_8):
    """One iteration of the series loop."""

    arg_9 = 2. * arg_4 / arg_1 + arg_5
    arg_10 = -2. * arg_4 / arg_1 ** 2 + arg_6
    arg_5 = 1. / arg_9
    arg_6 = -arg_10 / arg_9 ** 2

    arg_11 = tf.sin(arg_4 * arg_0) / arg_4 + arg_7
    arg_7 = arg_5 * arg_11
    arg_8 = (arg_6 * arg_11 +
                          arg_5 * arg_8)
    arg_4 -= 1.

    return arg_4, arg_5, arg_6, arg_7, arg_8

  (arg_12, arg_12, arg_12, arg_7, arg_8) = tf.while_loop(
      cond=lambda arg_4, *arg_12: arg_4 > 0.,
      body=loop_body,
      loop_vars=(
          arg_2,  # n
          tf.zeros_like(arg_0, name="rn"),
          tf.zeros_like(arg_0, name="drn_dconcentration"),
          tf.zeros_like(arg_0, name="vn"),
          tf.zeros_like(arg_0, name="dvn_dconcentration"),
      ),
  )

  arg_13 = .5 + arg_0 / (2. * np.pi) + arg_7 / np.pi
  arg_14 = arg_8 / np.pi

  # Clip the result to [0, 1].
  arg_15 = tf.clip_by_value(arg_13, 0., 1.)
  # The clipped values do not depend on concentration anymore, so set their
  # derivative to zero.
  arg_14 *= tf.cast((arg_13 >= 0.) & (arg_13 <= 1.), arg_3)

  return arg_15, arg_14


def Func(arg_0, arg_1, arg_2):
  """Computes the von Mises CDF and its derivative via Normal approximation."""

  def cdf_func(arg_1):
    """A helper function that is passed to value_and_gradient."""
    # z is an "almost Normally distributed" random variable.
    arg_3 = ((np.sqrt(2. / np.pi) / tf.math.bessel_i0e(arg_1)) *
         tf.sin(.5 * arg_0))

    # This is the correction described in [1] which reduces the error
    # of the Normal approximation.
    arg_4 = arg_3 ** 2
    arg_5 = arg_4 * arg_3
    arg_6 = arg_4 ** 2
    arg_7 = 24. * arg_1
    arg_8 = 56.

    arg_9 = arg_3 - arg_5 / ((arg_7 - 2. * arg_4 - 16.) / 3. -
                   (arg_6 + (7. / 4.) * arg_4 + 167. / 2.) / (arg_7 - arg_8 - arg_4 + 3.)) ** 2

    arg_10 = normal.Normal(tf.cast(0., arg_2), tf.cast(1., arg_2))

    return arg_10.cdf(arg_9)

  return value_and_gradient(cdf_func, arg_1)


def Func(
    arg_0,
    arg_1,
    arg_2=None,
    arg_3=0.5,
    arg_4=0.9,
    arg_5=None,
    arg_6=None):
  """Performs one step of the differential evolution algorithm.

  Args:
    objective_function:  A Python callable that accepts a batch of possible
      solutions and returns the values of the objective function at those
      arguments as a rank 1 real `Tensor`. This specifies the function to be
      minimized. The input to this callable may be either a single `Tensor`
      or a Python `list` of `Tensor`s. The signature must match the format of
      the argument `population`. (i.e. objective_function(*population) must
      return the value of the function to be minimized).
    population:  `Tensor` or Python `list` of `Tensor`s representing the
      current population vectors. Each `Tensor` must be of the same real dtype.
      The first dimension indexes individual population members while the
      rest of the dimensions are consumed by the value function. For example,
      if the population is a single `Tensor` of shape [n, m1, m2], then `n` is
      the population size and the output of `objective_function` applied to the
      population is a `Tensor` of shape [n]. If the population is a python
      list of `Tensor`s then each `Tensor` in the list should have the first
      axis of a common size, say `n` and `objective_function(*population)`
      should return a `Tensor of shape [n]. The population must have at least
      4 members for the algorithm to work correctly.
    population_values: A `Tensor` of rank 1 and real dtype. The result of
      applying `objective_function` to the `population`. If not supplied it is
      computed using the `objective_function`.
      Default value: None.
    differential_weight: Real scalar `Tensor`. Must be positive and less than
      2.0. The parameter controlling the strength of mutation.
      Default value: 0.5
    crossover_prob: Real scalar `Tensor`. Must be between 0 and 1. The
      probability of recombination per site.
      Default value: 0.9
    seed: `int` or None. The random seed for this `Op`. If `None`, no seed is
      applied.
      Default value: None.
    name: (Optional) Python str. The name prefixed to the ops created by this
      function. If not supplied, the default name 'Func' is
      used.
      Default value: None

  Returns:
    A sequence containing the following elements (in order):
    next_population: A `Tensor` or Python `list` of `Tensor`s of the same
      structure as the input population. The population at the next generation.
    next_population_values: A `Tensor` of same shape and dtype as input
      `population_values`. The function values for the `next_population`.
  """
  with tf.compat.v1.name_scope(
      arg_6, 'Func',
      [arg_1, arg_2, arg_3, arg_4]):
    arg_1, arg_7 = _ensure_list(arg_1)
    if arg_2 is None:
      arg_2 = arg_0(*arg_1)
    arg_8 = tf.shape(input=arg_1[0])[0]
    arg_9 = distributions.SeedStream(arg_5, salt='Func')
    arg_10 = _get_mixing_indices(arg_8, arg_5=arg_9())
    # Construct the mutated solution vectors. There is one for each member of
    # the population.
    arg_11 = _get_mutants(arg_1,
                           arg_8,
                           arg_10,
                           arg_3)
    # Perform recombination between the parents and the mutants.
    arg_12 = _binary_crossover(arg_1,
                                   arg_8,
                                   arg_11,
                                   arg_4,
                                   arg_5=arg_9())
    arg_13 = arg_0(*arg_12)
    if arg_2 is None:
      arg_2 = arg_0(*arg_1)

    arg_14 = tf.zeros_like(arg_2) + np.inf

    arg_2 = tf.where(
        tf.math.is_nan(arg_2), x=arg_14, y=arg_2)

    arg_15 = arg_13 < arg_2
    arg_16 = [
        tf.where(arg_15, x=candidates_part, y=population_part)
        for candidates_part, population_part in zip(arg_12, arg_1)
    ]
    arg_17 = tf.where(arg_15, x=arg_13, y=arg_2)

  return arg_16, arg_17


def Func(arg_0,
                      arg_1,
                      arg_2,
                      arg_3,
                      arg_4,
                      arg_5,
                      arg_6,
                      arg_7,
                      arg_8,
                      arg_9,
                      arg_10):
  """Processes initial args."""
  arg_11 = False
  if arg_2 is not None:
    arg_2, arg_11 = _ensure_list(arg_2)

  if arg_1 is not None:
    arg_1, arg_11 = _ensure_list(arg_1)

  arg_12 = _get_starting_population(arg_1,
                                        arg_2,
                                        arg_3,
                                        arg_4,
                                        arg_10=arg_10)

  arg_8 = tf.convert_to_tensor(
      value=arg_8, dtype=arg_12[0].dtype.base_dtype)

  arg_9 = tf.convert_to_tensor(value=arg_9)
  arg_13 = arg_0(*arg_12)
  if arg_5 is not None:
    arg_5 = tf.convert_to_tensor(value=arg_5)
  arg_6 = tf.convert_to_tensor(
      value=arg_6, dtype=arg_13.dtype.base_dtype)
  arg_7 = tf.convert_to_tensor(
      value=arg_7, dtype=arg_12[0].dtype.base_dtype)
  return (arg_11,
          arg_12,
          arg_13,
          arg_5,
          arg_6,
          arg_7,
          arg_8,
          arg_9)


def Func(arg_0, arg_1):
  """Finds the population member with the lowest value."""
  arg_2 = tf.math.reduce_min(input_tensor=arg_1)
  arg_3 = tf.where(tf.math.equal(arg_1, arg_2))[0, 0]

  return ([arg_4[arg_3] for arg_4 in arg_0],
          arg_2)


def Func(arg_0,
                 arg_1,
                 arg_2,
                 arg_3):
  """Computes the mutatated vectors for each population member.

  Args:
    population:  Python `list` of `Tensor`s representing the
      current population vectors. Each `Tensor` must be of the same real dtype.
      The first dimension of each `Tensor` indexes individual
      population members. For example, if the population is a list with a
      single `Tensor` of shape [n, m1, m2], then `n` is the population size and
      the shape of an individual solution is [m1, m2].
      If there is more than one element in the population, then each `Tensor`
      in the list should have the first axis of the same size.
    population_size: Scalar integer `Tensor`. The size of the population.
    mixing_indices: `Tensor` of integral dtype and shape [n, 3] where `n` is the
      number of members in the population. Each element of the `Tensor` must be
      a valid index into the first dimension of the population (i.e range
      between `0` and `n-1` inclusive).
    differential_weight: Real scalar `Tensor`. Must be positive and less than
      2.0. The parameter controlling the strength of mutation.

  Returns:
    mutants: `Tensor` or Python `list` of `Tensor`s of the same shape and dtype
      as the input population. The mutated vectors.
  """
  arg_2 = tf.reshape(arg_2, [-1])
  arg_4 = tf.stack([1.0, arg_3, -arg_3])
  def _mutant_part(arg_5):
    arg_6 = tf.gather(arg_5, arg_2)
    arg_6 = tf.transpose(
        a=tf.reshape(arg_6, [arg_1, 3, -1]), perm=[0, 2, 1])
    return tf.math.reduce_sum(input_tensor=arg_6 * arg_4, axis=-1)

  return [_mutant_part(arg_5) for arg_5 in arg_0]


def Func(arg_0, arg_1=None, arg_2=None):
  """Generates an array of indices suitable for mutation operation.

  The mutation operation in differential evolution requires that for every
  element of the population, three distinct other elements be chosen to produce
  a trial candidate. This function generates an array of shape [size, 3]
  satisfying the properties that:
    (a). array[i, :] does not contain the index 'i'.
    (b). array[i, :] does not contain any overlapping indices.
    (c). All elements in the array are between 0 and size - 1 inclusive.

  Args:
    size: Scalar integer `Tensor`. The number of samples as well as a the range
      of the indices to sample from.
    seed: `int` or None. The random seed for this `Op`. If `None`, no seed is
      applied.
      Default value: `None`.
    name:  Python `str` name prefixed to Ops created by this function.
      Default value: 'get_mixing_indices'.

  Returns:
    sample: A `Tensor` of shape [size, 3] and same dtype as `size` containing
    samples without replacement between 0 and size - 1 (inclusive) with the
    `i`th row not including the number `i`.
  """
  with tf.compat.v1.name_scope(
      arg_2, default_name='get_mixing_indices', values=[arg_0]):
    arg_0 = tf.convert_to_tensor(value=arg_0)
    arg_3 = arg_0.dtype
    arg_4 = distributions.SeedStream(arg_1, salt='get_mixing_indices')
    arg_5 = tf.random.uniform([arg_0],
                              maxval=arg_0-1,
                              arg_3=arg_3,
                              arg_1=arg_4())
    arg_6 = tf.random.uniform([arg_0],
                               maxval=arg_0-2,
                               arg_3=arg_3,
                               arg_1=arg_4())
    arg_7 = tf.random.uniform([arg_0],
                              maxval=arg_0-3,
                              arg_3=arg_3,
                              arg_1=arg_4())

    # Shift second if it is on top of or to the right of first
    arg_6 = tf.where(arg_5 < arg_6, x=arg_6, y=arg_6 + 1)
    arg_8 = tf.math.minimum(arg_5, arg_6)
    arg_9 = tf.math.maximum(arg_5, arg_6)
    # Shift the third one so it does not coincide with either the first or the
    # second number. Assuming first < second, shift by 1 if the number is in
    # [first, second) and by 2 if the number is greater than or equal to the
    # second.
    arg_7 = tf.where(arg_7 < arg_8, x=arg_7, y=arg_7 + 1)
    arg_7 = tf.where(arg_7 < arg_9, x=arg_7, y=arg_7 + 1)
    arg_10 = tf.stack([arg_5, arg_6, arg_7], axis=1)
    arg_11 = tf.expand_dims(tf.range(arg_0), axis=-1)
    arg_10 = tf.where(arg_10 < arg_11, x=arg_10, y=arg_10 + 1)
    return arg_10


def Func(arg_0):
  """Converts the input arg to a list if it is not a list already.

  Args:
    tensor_or_list: A `Tensor` or a Python list of `Tensor`s. The argument to
      convert to a list of `Tensor`s.

  Returns:
    A tuple of two elements. The first is a Python list of `Tensor`s containing
    the original arguments. The second is a boolean indicating whether
    the original argument was a list or tuple already.
  """
  if isinstance(arg_0, (list, tuple)):
    return list(arg_0), True
  return [arg_0], False


def Func(arg_0, arg_1, arg_2):
  """Gets a Tensor of type `dtype`, 0 if `tol` is None, validation optional."""
  if arg_0 is None:
    return tf.convert_to_tensor(value=0, arg_1=arg_1)

  arg_0 = tf.convert_to_tensor(value=arg_0, arg_1=arg_1)
  if arg_2:
    arg_0 = distribution_util.with_dependencies([
        assert_util.assert_non_negative(
            arg_0, message="Argument 'tol' must be non-negative")
    ], arg_0)
  return arg_0


def Func(arg_0, arg_1, arg_2=None):
  """Soft Thresholding operator.

  This operator is defined by the equations

  ```none
                                { x[i] - gamma,  x[i] >   gamma
  SoftThreshold(x, gamma)[i] =  { 0,             x[i] ==  gamma
                                { x[i] + gamma,  x[i] <  -gamma
  ```

  In the context of proximal gradient methods, we have

  ```none
  SoftThreshold(x, gamma) = prox_{gamma L1}(x)
  ```

  where `prox` is the proximity operator.  Thus the soft thresholding operator
  is used in proximal gradient descent for optimizing a smooth function with
  (non-smooth) L1 regularization, as outlined below.

  The proximity operator is defined as:

  ```none
  prox_r(x) = argmin{ r(z) + 0.5 ||x - z||_2**2 : z },
  ```

  where `r` is a (weakly) convex function, not necessarily differentiable.
  Because the L2 norm is strictly convex, the above argmin is unique.

  One important application of the proximity operator is as follows.  Let `L` be
  a convex and differentiable function with Lipschitz-continuous gradient.  Let
  `R` be a convex lower semicontinuous function which is possibly
  nondifferentiable.  Let `gamma` be an arbitrary positive real.  Then

  ```none
  x_star = argmin{ L(x) + R(x) : x }
  ```

  if and only if the fixed-point equation is satisfied:

  ```none
  x_star = prox_{gamma R}(x_star - gamma grad L(x_star))
  ```

  Proximal gradient descent thus typically consists of choosing an initial value
  `x^{(0)}` and repeatedly applying the update

  ```none
  x^{(k+1)} = prox_{gamma^{(k)} R}(x^{(k)} - gamma^{(k)} grad L(x^{(k)}))
  ```

  where `gamma` is allowed to vary from iteration to iteration.  Specializing to
  the case where `R(x) = ||x||_1`, we minimize `L(x) + ||x||_1` by repeatedly
  applying the update

  ```
  x^{(k+1)} = SoftThreshold(x - gamma grad L(x^{(k)}), gamma)
  ```

  (This idea can also be extended to second-order approximations, although the
  multivariate case does not have a known closed form like above.)

  Args:
    x: `float` `Tensor` representing the input to the SoftThreshold function.
    threshold: nonnegative scalar, `float` `Tensor` representing the radius of
      the interval on which each coordinate of SoftThreshold takes the value
      zero.  Denoted `gamma` above.
    name: Python string indicating the name of the TensorFlow operation.
      Default value: `'Func'`.

  Returns:
    softthreshold: `float` `Tensor` with the same shape and dtype as `x`,
      representing the value of the SoftThreshold function.

  #### References

  [1]: Yu, Yao-Liang. The Proximity Operator.
       https://www.cs.cmu.edu/~suvrit/teach/yaoliang_proximity.pdf

  [2]: Wikipedia Contributors. Proximal gradient methods for learning.
       _Wikipedia, The Free Encyclopedia_, 2018.
       https://en.wikipedia.org/wiki/Proximal_gradient_methods_for_learning

  """
  # https://math.stackexchange.com/questions/471339/derivation-of-soft-thresholding-operator
  with tf.compat.v1.name_scope(arg_2, 'Func', [arg_0, arg_1]):
    arg_0 = tf.convert_to_tensor(value=arg_0, arg_2='x')
    arg_1 = tf.convert_to_tensor(
        value=arg_1, dtype=arg_0.dtype, arg_2='threshold')
    return tf.sign(arg_0) * tf.maximum(tf.abs(arg_0) - arg_1, 0.)


def Func(arg_0, arg_1, arg_2,
                                    arg_3=None):
  """Clips values to a specified min and max while leaving gradient unaltered.

  Like `tf.clip_by_value`, this function returns a tensor of the same type and
  shape as input `t` but with values clamped to be no smaller than to
  `clip_value_min` and no larger than `clip_value_max`. Unlike
  `tf.clip_by_value`, the gradient is unaffected by this op, i.e.,

  ```python
  tf.gradients(tfp.math.Func(x), x)[0]
  # ==> ones_like(x)
  ```

  Note: `clip_value_min` needs to be smaller or equal to `clip_value_max` for
  correct results.

  Args:
    t: A `Tensor`.
    clip_value_min: A scalar `Tensor`, or a `Tensor` with the same shape
      as `t`. The minimum value to clip by.
    clip_value_max: A scalar `Tensor`, or a `Tensor` with the same shape
      as `t`. The maximum value to clip by.
    name: A name for the operation (optional).
      Default value: `'Func'`.

  Returns:
    clipped_t: A clipped `Tensor`.
  """
  with tf.compat.v1.name_scope(arg_3, 'Func',
                               [arg_0, arg_1, arg_2]):
    arg_0 = tf.convert_to_tensor(value=arg_0, arg_3='t')
    arg_4 = tf.clip_by_value(arg_0, arg_1, arg_2)
    return arg_0 + tf.stop_gradient(arg_4 - arg_0)


def Func(arg_0, arg_1):
  """Build an iterator over training batches."""

  arg_2 = tf.data.Dataset.from_tensor_slices(arg_0)
  arg_3 = arg_2.shuffle(
      50000, reshuffle_each_iteration=True).repeat().batch(arg_1)
  arg_4 = tf.compat.v1.data.make_one_shot_iterator(arg_3)
  arg_5 = arg_4.get_next()
  return arg_5


def Func(arg_0, arg_1):
    """Converts a sequence of productions into a string of terminal symbols.

    Args:
      productions: Tensor of shape [1, num_productions, num_production_rules].
        Slices along the `num_productions` dimension represent one-hot vectors.

    Returns:
      str that concatenates all terminal symbols from `productions`.

    Raises:
      ValueError: If the first production rule does not begin with
        `self.start_symbol`.
    """
    arg_2 = []
    for arg_3 in tf.unstack(arg_1, axis=1):
      arg_4, arg_5 = arg_0.production_rules[tf.argmax(input=arg_3, axis=-1)]
      if not arg_2:  # first iteration
        if arg_4 != arg_0.start_symbol:
          raise ValueError("`productions` must begin with `self.start_symbol`.")
        arg_2 = arg_5
      else:
        # Greedily unroll the nonterminal symbols based on the first occurrence
        # in a linear sequence.
        arg_6 = arg_2.index(arg_4)
        arg_2 = arg_2[:arg_6] + arg_5 + arg_2[arg_6 + 1:]
    arg_7 = "".join(arg_2)
    return arg_7


def Func(arg_0, arg_1):
    """Runs the model forward to generate a sequence of productions.

    Args:
      inputs: Unused.

    Returns:
      productions: Tensor of shape [1, num_productions, num_production_rules].
        Slices along the `num_productions` dimension represent one-hot vectors.
    """
    del arg_1  # unused
    arg_2 = ed.MultivariateNormalDiag(loc=tf.zeros(arg_0.latent_size),
                                            sample_shape=1,
                                            name="latent_code")
    arg_3 = arg_0.lstm.zero_state(1, dtype=tf.float32)
    arg_4 = 0
    arg_5 = []
    arg_6 = [arg_0.grammar.start_symbol]
    while arg_6:
      arg_7 = arg_6.pop()
      arg_8, arg_3 = arg_0.lstm(arg_2, arg_3)
      arg_9 = (arg_0.output_layer(arg_8) +
                arg_0.grammar.mask(arg_7, on_value=0., off_value=-1e9))
      arg_10 = ed.OneHotCategorical(arg_9=arg_9,
                                        name="production_" + str(arg_4))
      arg_11, arg_12 = arg_0.grammar.production_rules[tf.argmax(
          input=arg_10, axis=-1)]
      for arg_7 in arg_12:
        if arg_7 in arg_0.grammar.nonterminal_symbols:
          arg_6.append(arg_7)
      arg_5.append(arg_10)
      arg_4 += 1
    return tf.stack(arg_5, axis=1)


def Func(arg_0, arg_1):
    """Runs the model forward to return a stochastic encoding.

    Args:
      inputs: Tensor of shape [1, num_productions, num_production_rules]. It is
        a sequence of productions of length `num_productions`. Each production
        is a one-hot vector of length `num_production_rules`: it determines
        which production rule the production corresponds to.

    Returns:
      latent_code_posterior: A random variable capturing a sample from the
        variational distribution, of shape [1, self.latent_size].
    """
    arg_2 = arg_0.encoder_net(tf.cast(arg_1, tf.float32))
    return ed.MultivariateNormalDiag(
        loc=arg_2[..., :arg_0.latent_size],
        scale_diag=tf.nn.softplus(arg_2[..., arg_0.latent_size:]),
        name="latent_code_posterior")


def Func(arg_0, arg_1):
    """Integral of the `hat` function, used for sampling.

    We choose a `hat` function, h(x) = x^(-power), which is a continuous
    (unnormalized) density touching each positive integer at the (unnormalized)
    pmf. This function implements `hat` integral: H(x) = int_x^inf h(t) dt;
    which is needed for sampling purposes.

    Arguments:
      x: A Tensor of points x at which to evaluate H(x).

    Returns:
      A Tensor containing evaluation H(x) at x.
    """
    arg_1 = tf.cast(arg_1, arg_0.power.dtype)
    arg_2 = arg_0.power - 1.
    return tf.exp((-arg_2) * tf.math.log1p(arg_1) - tf.math.log(arg_2))


def Func(arg_0, arg_1):
    """Inverse function of _hat_integral."""
    arg_1 = tf.cast(arg_1, arg_0.power.dtype)
    arg_2 = arg_0.power - 1.
    return tf.math.expm1(-(tf.math.log(arg_2) + tf.math.log(arg_1)) / arg_2)


def Func(arg_0, arg_1=None, arg_2=False, arg_3=None):
  """Compute the matrix rank; the number of non-zero SVD singular values.

  Arguments:
    a: (Batch of) `float`-like matrix-shaped `Tensor`(s) which are to be
      pseudo-inverted.
    tol: Threshold below which the singular value is counted as "zero".
      Default value: `None` (i.e., `eps * max(rows, cols) * max(singular_val)`).
    validate_args: When `True`, additional assertions might be embedded in the
      graph.
      Default value: `False` (i.e., no graph assertions are added).
    name: Python `str` prefixed to ops created by this function.
      Default value: "Func".

  Returns:
    Func: (Batch of) `int32` scalars representing the number of non-zero
      singular values.
  """
  with tf.compat.v1.name_scope(arg_3, 'Func', [arg_0, arg_1]):
    arg_0 = tf.convert_to_tensor(value=arg_0, dtype_hint=tf.float32, arg_3='a')
    arg_4 = _maybe_validate_matrix(arg_0, arg_2)
    if arg_4:
      with tf.control_dependencies(arg_4):
        arg_0 = tf.identity(arg_0)
    arg_5 = tf.linalg.svd(arg_0, compute_uv=False)
    if arg_1 is None:
      if arg_0.shape[-2:].is_fully_defined():
        arg_6 = np.max(arg_0.shape[-2:].as_list())
      else:
        arg_6 = tf.reduce_max(input_tensor=tf.shape(input=arg_0)[-2:])
      arg_7 = np.finfo(arg_0.dtype.as_numpy_dtype).eps
      arg_1 = (arg_7 * tf.cast(arg_6, arg_0.dtype) *
             tf.reduce_max(input_tensor=arg_5, axis=-1, keepdims=True))
    return tf.reduce_sum(input_tensor=tf.cast(arg_5 > arg_1, tf.int32), axis=-1)


def Func(arg_0, arg_1=None, arg_2=False, arg_3=None):
  """Compute the Moore-Penrose pseudo-inverse of a matrix.

  Calculate the [generalized inverse of a matrix](
  https://en.wikipedia.org/wiki/Moore%E2%80%93Penrose_inverse) using its
  singular-value decomposition (SVD) and including all large singular values.

  The pseudo-inverse of a matrix `A`, is defined as: "the matrix that 'solves'
  [the least-squares problem] `A @ x = b`," i.e., if `x_hat` is a solution, then
  `A_Func` is the matrix such that `x_hat = A_Func @ b`. It can be shown that if
  `U @ Sigma @ V.T = A` is the singular value decomposition of `A`, then
  `A_Func = V @ inv(Sigma) U^T`. [(Strang, 1980)][1]

  This function is analogous to [`numpy.linalg.Func`](
  https://docs.scipy.org/doc/numpy/reference/generated/numpy.linalg.Func.html).
  It differs only in default value of `rcond`. In `numpy.linalg.Func`, the
  default `rcond` is `1e-15`. Here the default is
  `10. * max(num_rows, num_cols) * np.finfo(dtype).eps`.

  Args:
    a: (Batch of) `float`-like matrix-shaped `Tensor`(s) which are to be
      pseudo-inverted.
    rcond: `Tensor` of small singular value cutoffs.  Singular values smaller
      (in modulus) than `rcond` * largest_singular_value (again, in modulus) are
      set to zero. Must broadcast against `tf.shape(a)[:-2]`.
      Default value: `10. * max(num_rows, num_cols) * np.finfo(a.dtype).eps`.
    validate_args: When `True`, additional assertions might be embedded in the
      graph.
      Default value: `False` (i.e., no graph assertions are added).
    name: Python `str` prefixed to ops created by this function.
      Default value: "Func".

  Returns:
    a_Func: The pseudo-inverse of input `a`. Has same shape as `a` except
      rightmost two dimensions are transposed.

  Raises:
    TypeError: if input `a` does not have `float`-like `dtype`.
    ValueError: if input `a` has fewer than 2 dimensions.

  #### Examples

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

  a = tf.constant([[1.,  0.4,  0.5],
                   [0.4, 0.2,  0.25],
                   [0.5, 0.25, 0.35]])
  tf.matmul(tfp.math.Func(a), a)
  # ==> array([[1., 0., 0.],
               [0., 1., 0.],
               [0., 0., 1.]], dtype=float32)

  a = tf.constant([[1.,  0.4,  0.5,  1.],
                   [0.4, 0.2,  0.25, 2.],
                   [0.5, 0.25, 0.35, 3.]])
  tf.matmul(tfp.math.Func(a), a)
  # ==> array([[ 0.76,  0.37,  0.21, -0.02],
               [ 0.37,  0.43, -0.33,  0.02],
               [ 0.21, -0.33,  0.81,  0.01],
               [-0.02,  0.02,  0.01,  1.  ]], dtype=float32)
  ```

  #### References

  [1]: G. Strang. "Linear Algebra and Its Applications, 2nd Ed." Academic Press,
       Inc., 1980, pp. 139-142.
  """
  with tf.compat.v1.name_scope(arg_3, 'Func', [arg_0, arg_1]):
    arg_0 = tf.convert_to_tensor(value=arg_0, arg_3='a')

    arg_4 = _maybe_validate_matrix(arg_0, arg_2)
    if arg_4:
      with tf.control_dependencies(arg_4):
        arg_0 = tf.identity(arg_0)

    arg_5 = arg_0.dtype.as_numpy_dtype

    if arg_1 is None:
      def get_dim_size(arg_6):
        if tf.compat.dimension_value(arg_0.shape[arg_6]) is not None:
          return tf.compat.dimension_value(arg_0.shape[arg_6])
        return tf.shape(input=arg_0)[arg_6]

      arg_7 = get_dim_size(-2)
      arg_8 = get_dim_size(-1)
      if isinstance(arg_7, int) and isinstance(arg_8, int):
        arg_9 = float(max(arg_7, arg_8))
      else:
        arg_9 = tf.cast(tf.maximum(arg_7, arg_8), arg_5)
      arg_1 = 10. * arg_9 * np.finfo(arg_5).eps

    arg_1 = tf.convert_to_tensor(value=arg_1, arg_5=arg_5, arg_3='rcond')

    # Calculate pseudo inverse via SVD.
    # Note: if a is symmetric then u == v. (We might observe additional
    # performance by explicitly setting `v = u` in such cases.)
    [
        arg_10,         # Sigma
        arg_11,   # U
        arg_12,  # V
    ] = tf.linalg.svd(arg_0, full_matrices=False, compute_uv=True)

    # Saturate small singular values to inf. This has the effect of make
    # `1. / s = 0.` while not resulting in `NaN` gradients.
    arg_13 = arg_1 * tf.reduce_max(input_tensor=arg_10, axis=-1)
    arg_10 = tf.where(
        arg_10 > arg_13[..., tf.newaxis], arg_10,
        tf.fill(tf.shape(input=arg_10), np.array(np.inf, arg_5)))

    # Although `a == tf.matmul(u, s * v, transpose_b=True)` we swap
    # `u` and `v` here so that `tf.matmul(Func(A), A) = tf.eye()`, i.e.,
    # a matrix inverse has "transposed" semantics.
    arg_14 = tf.matmul(
        arg_12 / arg_10[..., tf.newaxis, :],
        arg_11,
        adjoint_b=True)

    if arg_0.shape.ndims is not None:
      arg_14.set_shape(arg_0.shape[:-2].concatenate([arg_0.shape[-1], arg_0.shape[-2]]))

    return arg_14


def Func(arg_0, arg_1, arg_2,
             arg_3=False,
             arg_4=None):
  """Solves systems of linear eqns `A X = RHS`, given LU factorizations.

  Note: this function does not verify the implied matrix is actually invertible
  nor is this condition checked even when `validate_args=True`.

  Args:
    lower_upper: `lu` as returned by `tf.linalg.lu`, i.e., if
      `matmul(P, matmul(L, U)) = X` then `lower_upper = L + U - eye`.
    perm: `p` as returned by `tf.linag.lu`, i.e., if
      `matmul(P, matmul(L, U)) = X` then `perm = argmax(P)`.
    rhs: Matrix-shaped float `Tensor` representing targets for which to solve;
      `A X = RHS`. To handle vector cases, use:
      `Func(..., rhs[..., tf.newaxis])[..., 0]`.
    validate_args: Python `bool` indicating whether arguments should be checked
      for correctness. Note: this function does not verify the implied matrix is
      actually invertible, even when `validate_args=True`.
      Default value: `False` (i.e., don't validate arguments).
    name: Python `str` name given to ops managed by this object.
      Default value: `None` (i.e., "Func").

  Returns:
    x: The `X` in `A @ X = RHS`.

  #### Examples

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

  x = [[[1., 2],
        [3, 4]],
       [[7, 8],
        [3, 4]]]
  inv_x = tfp.math.Func(*tf.linalg.lu(x), rhs=tf.eye(2))
  tf.assert_near(tf.matrix_inverse(x), inv_x)
  # ==> True
  ```

  """

  with tf.compat.v1.name_scope(arg_4, 'Func', [arg_0, arg_1, arg_2]):
    arg_0 = tf.convert_to_tensor(
        value=arg_0, dtype_hint=tf.float32, arg_4='lower_upper')
    arg_1 = tf.convert_to_tensor(value=arg_1, dtype_hint=tf.int32, arg_4='perm')
    arg_2 = tf.convert_to_tensor(
        value=arg_2, dtype_hint=arg_0.dtype, arg_4='rhs')

    arg_5 = _Func_assertions(arg_0, arg_1, arg_2, arg_3)
    if arg_5:
      with tf.control_dependencies(arg_5):
        arg_0 = tf.identity(arg_0)
        arg_1 = tf.identity(arg_1)
        arg_2 = tf.identity(arg_2)

    if arg_2.shape.ndims == 2 and arg_1.shape.ndims == 1:
      # Both rhs and perm have scalar batch_shape.
      arg_6 = tf.gather(arg_2, arg_1, axis=-2)
    else:
      # Either rhs or perm have non-scalar batch_shape or we can't determine
      # this information statically.
      arg_7 = tf.shape(input=arg_2)
      arg_8 = tf.broadcast_dynamic_shape(
          arg_7[:-2],
          tf.shape(input=arg_1)[:-1])
      arg_9, arg_10 = arg_7[-2], arg_7[-1]
      arg_11 = tf.concat([arg_8, [arg_9, arg_10]], axis=0)

      # Tile out rhs.
      arg_12 = tf.broadcast_to(arg_2, arg_11)
      arg_12 = tf.reshape(arg_12, [-1, arg_9, arg_10])

      # Tile out perm and add batch indices.
      arg_13 = tf.broadcast_to(arg_1, arg_11[:-1])
      arg_13 = tf.reshape(arg_13, [-1, arg_9])
      arg_14 = tf.reduce_prod(input_tensor=arg_8)
      arg_15 = tf.broadcast_to(
          tf.range(arg_14)[:, tf.newaxis],
          [arg_14, arg_9])
      arg_13 = tf.stack([arg_15, arg_13],
                                axis=-1)

      arg_6 = tf.gather_nd(arg_12, arg_13)
      arg_6 = tf.reshape(arg_6, arg_11)

    arg_16 = tf.linalg.set_diag(
        tf.linalg.band_part(arg_0, num_lower=-1, num_upper=0),
        tf.ones(tf.shape(input=arg_0)[:-1], dtype=arg_0.dtype))
    return linear_operator_util.matrix_triangular_solve_with_broadcast(
        arg_0,  # Only upper is accessed.
        linear_operator_util.matrix_triangular_solve_with_broadcast(
            arg_16, arg_6),
        arg_16=False)


def Func(arg_0, arg_1, arg_2=False, arg_3=None):
  """Computes a matrix inverse given the matrix's LU decomposition.

  This op is conceptually identical to,

  ````python
  inv_X = tf.Func(*tf.linalg.lu(X))
  tf.assert_near(tf.matrix_inverse(X), inv_X)
  # ==> True
  ```

  Note: this function does not verify the implied matrix is actually invertible
  nor is this condition checked even when `validate_args=True`.

  Args:
    lower_upper: `lu` as returned by `tf.linalg.lu`, i.e., if
      `matmul(P, matmul(L, U)) = X` then `lower_upper = L + U - eye`.
    perm: `p` as returned by `tf.linag.lu`, i.e., if
      `matmul(P, matmul(L, U)) = X` then `perm = argmax(P)`.
    validate_args: Python `bool` indicating whether arguments should be checked
      for correctness. Note: this function does not verify the implied matrix is
      actually invertible, even when `validate_args=True`.
      Default value: `False` (i.e., don't validate arguments).
    name: Python `str` name given to ops managed by this object.
      Default value: `None` (i.e., "Func").

  Returns:
    inv_x: The matrix_inv, i.e.,
      `tf.matrix_inverse(tfp.math.lu_reconstruct(lu, perm))`.

  #### Examples

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

  x = [[[3., 4], [1, 2]],
       [[7., 8], [3, 4]]]
  inv_x = tfp.math.Func(*tf.linalg.lu(x))
  tf.assert_near(tf.matrix_inverse(x), inv_x)
  # ==> True
  ```

  """

  with tf.compat.v1.name_scope(arg_3, 'Func', [arg_0, arg_1]):
    arg_0 = tf.convert_to_tensor(
        value=arg_0, dtype_hint=tf.float32, arg_3='lower_upper')
    arg_1 = tf.convert_to_tensor(value=arg_1, dtype_hint=tf.int32, arg_3='perm')
    arg_4 = _lu_reconstruct_assertions(arg_0, arg_1, arg_2)
    if arg_4:
      with tf.control_dependencies(arg_4):
        arg_0 = tf.identity(arg_0)
        arg_1 = tf.identity(arg_1)
    arg_5 = tf.shape(input=arg_0)
    return lu_solve(
        arg_0, arg_1,
        rhs=tf.eye(arg_5[-1], batch_shape=arg_5[:-2], dtype=arg_0.dtype),
        arg_2=False)


def Func(arg_0, arg_1, arg_2, arg_3):
  """Returns list of assertions related to `lu_solve` assumptions."""
  arg_4 = _lu_reconstruct_assertions(arg_0, arg_1, arg_3)

  arg_5 = 'Input `rhs` must have at least 2 dimensions.'
  if arg_2.shape.ndims is not None:
    if arg_2.shape.ndims < 2:
      raise ValueError(arg_5)
  elif arg_3:
    arg_4.append(
        tf.compat.v1.assert_rank_at_least(arg_2, rank=2, arg_5=arg_5))

  arg_5 = '`lower_upper.shape[-1]` must equal `rhs.shape[-1]`.'
  if (tf.compat.dimension_value(arg_0.shape[-1]) is not None and
      tf.compat.dimension_value(arg_2.shape[-2]) is not None):
    if arg_0.shape[-1] != arg_2.shape[-2]:
      raise ValueError(arg_5)
  elif arg_3:
    arg_4.append(
        tf.compat.v1.assert_equal(
            tf.shape(input=arg_0)[-1],
            tf.shape(input=arg_2)[-2],
            arg_5=arg_5))

  return arg_4


def Func(arg_0):
  """Returns a block diagonal rank 2 SparseTensor from a batch of SparseTensors.

  Args:
    sp_a: A rank 3 `SparseTensor` representing a batch of matrices.

  Returns:
    sp_block_diag_a: matrix-shaped, `float` `SparseTensor` with the same dtype
    as `sparse_or_matrix`, of shape [B * M, B * N] where `sp_a` has shape
    [B, M, N]. Each [M, N] batch of `sp_a` is lined up along the diagonal.
  """
  # Construct the matrix [[M, N], [1, 0], [0, 1]] which would map the index
  # (b, i, j) to (Mb + i, Nb + j). This effectively creates a block-diagonal
  # matrix of dense shape [B * M, B * N].
  # Note that this transformation doesn't increase the number of non-zero
  # entries in the SparseTensor.
  arg_1 = tf.convert_to_tensor(value=_get_shape(arg_0, tf.int64))
  arg_2 = tf.concat([[arg_1[-2:]], tf.eye(2, dtype=tf.int64)], axis=0)
  arg_3 = tf.matmul(arg_0.indices, arg_2)
  arg_4 = arg_1[0] * arg_1[1:]
  return tf.SparseTensor(
      arg_3=arg_3, values=arg_0.values, arg_4=arg_4)


def Func(arg_0, arg_1, arg_2,
                                     arg_3):
  """Computes the neg-log-likelihood gradient and Fisher information for a GLM.

  Note that Fisher information is related to the Hessian of the log-likelihood
  by the equation

  ```none
  FisherInfo = E[Hessian with respect to model_coefficients of -LogLikelihood(
      Y | model_matrix, model_coefficients)]
  ```

  where `LogLikelihood` is the log-likelihood of a generalized linear model
  parameterized by `model_matrix` and `model_coefficients`, and the expectation
  is taken over Y, distributed according to the same GLM with the same parameter
  values.

  Args:
    model_matrix: (Batch of) matrix-shaped, `float` `Tensor` or `SparseTensor`
      where each row represents a sample's features.  Has shape `[N, n]` where
      `N` is the number of data samples and `n` is the number of features per
      sample.
    linear_response: (Batch of) vector-shaped `Tensor` with the same dtype as
      `model_matrix`, equal to `model_matix @ model_coefficients` where
      `model_coefficients` are the coefficients of the linear component of the
      GLM.
    response: (Batch of) vector-shaped `Tensor` with the same dtype as
      `model_matrix` where each element represents a sample's observed response
      (to the corresponding row of features).
    model: `tfp.glm.ExponentialFamily`-like instance, which specifies the link
      function and distribution of the GLM, and thus characterizes the negative
      log-likelihood. Must have sufficient statistic equal to the response, that
      is, `T(y) = y`.

  Returns:
    grad_neg_log_likelihood: (Batch of) vector-shaped `Tensor` with the same
      shape and dtype as a single row of `model_matrix`, representing the
      gradient of the negative log likelihood of `response` given linear
      response `linear_response`.
    fim_middle: (Batch of) vector-shaped `Tensor` with the same shape and dtype
      as a single column of `model_matrix`, satisfying the equation
      `Fisher information =
      Transpose(model_matrix)
      @ diag(fim_middle)
      @ model_matrix`.
  """
  # TODO(b/111926503): Determine whether there are some practical cases where it
  # is computationally favorable to compute the full FIM.
  arg_4, arg_5, arg_6 = arg_3(arg_1)

  arg_7 = (
      tf.math.is_finite(arg_6) & tf.not_equal(arg_6, 0.)
      & tf.math.is_finite(arg_5) & (arg_5 > 0.))

  def _mask_if_invalid(arg_8, arg_9):
    arg_9 = tf.fill(
        tf.shape(input=arg_8), value=np.array(arg_9, arg_8.dtype.as_numpy_dtype))
    return tf.where(arg_7, arg_8, arg_9)

  # TODO(b/111923449): Link to derivation once it's available.
  arg_10 = (arg_2 - arg_4) * _mask_if_invalid(arg_6, 1) / _mask_if_invalid(
      arg_5, np.inf)
  arg_11 = sparse_or_dense_matvecmul(
      arg_0, arg_10, adjoint_a=True)
  arg_12 = _mask_if_invalid(arg_6, 0.)**2 / _mask_if_invalid(
      arg_5, np.inf)
  return -arg_11, arg_12


def Func(arg_0,
               arg_1,
               arg_2,
               arg_3,
               arg_4,
               arg_5,
               arg_6=None,
               arg_7=None,
               arg_8=1,
               arg_9=None,
               arg_10=None):
  r"""Fits a GLM using coordinate-wise FIM-informed proximal gradient descent.

  This function uses a L1- and L2-regularized, second-order quasi-Newton method
  to find maximum-likelihood parameters for the given model and observed data.
  The second-order approximations use negative Fisher information in place of
  the Hessian, that is,

  ```none
  FisherInfo = E_Y[Hessian with respect to model_coefficients of -LogLikelihood(
      Y | model_matrix, current value of model_coefficients)]
  ```

  For large, sparse data sets, `model_matrix` should be supplied as a
  `SparseTensor`.

  Args:
    model_matrix: (Batch of) matrix-shaped, `float` `Tensor` or `SparseTensor`
      where each row represents a sample's features.  Has shape `[N, n]` where
      `N` is the number of data samples and `n` is the number of features per
      sample.
    response: (Batch of) vector-shaped `Tensor` with the same dtype as
      `model_matrix` where each element represents a sample's observed response
      (to the corresponding row of features).
    model: `tfp.glm.ExponentialFamily`-like instance, which specifies the link
      function and distribution of the GLM, and thus characterizes the negative
      log-likelihood which will be minimized. Must have sufficient statistic
      equal to the response, that is, `T(y) = y`.
    model_coefficients_start: (Batch of) vector-shaped, `float` `Tensor` with
      the same dtype as `model_matrix`, representing the initial values of the
      coefficients for the GLM regression.  Has shape `[n]` where `model_matrix`
      has shape `[N, n]`.
    tolerance: scalar, `float` `Tensor` representing the tolerance for each
      optiization step; see the `tolerance` argument of `Func_one_step`.
    l1_regularizer: scalar, `float` `Tensor` representing the weight of the L1
      regularization term.
    l2_regularizer: scalar, `float` `Tensor` representing the weight of the L2
      regularization term.
      Default value: `None` (i.e., no L2 regularization).
    maximum_iterations: Python integer specifying maximum number of iterations
      of the outer loop of the optimizer (i.e., maximum number of calls to
      `Func_one_step`).  After this many iterations of the outer loop, the
      algorithm will terminate even if the return value `model_coefficients` has
      not converged.
      Default value: `1`.
    maximum_full_sweeps_per_iteration: Python integer specifying the maximum
      number of coordinate descent sweeps allowed in each iteration.
      Default value: `1`.
    learning_rate: scalar, `float` `Tensor` representing a multiplicative factor
      used to dampen the proximal gradient descent steps.
      Default value: `None` (i.e., factor is conceptually `1`).
    name: Python string representing the name of the TensorFlow operation.
      The default name is `"Func"`.

  Returns:
    model_coefficients: (Batch of) `Tensor` of the same shape and dtype as
      `model_coefficients_start`, representing the computed model coefficients
      which minimize the regularized negative log-likelihood.
    is_converged: scalar, `bool` `Tensor` indicating whether the minimization
      procedure converged across all batches within the specified number of
      iterations.  Here convergence means that an iteration of the inner loop
      (`Func_one_step`) returns `True` for its `is_converged` output
      value.
    iter: scalar, `int` `Tensor` indicating the actual number of iterations of
      the outer loop of the optimizer completed (i.e., number of calls to
      `Func_one_step` before achieving convergence).

  #### Example

  ```python
  from __future__ import print_function
  import numpy as np
  import tensorflow as tf
  import tensorflow_probability as tfp
  tfd = tfp.distributions

  def make_dataset(n, d, link, scale=1., dtype=np.float32):
    model_coefficients = tfd.Uniform(
        low=np.array(-1, dtype), high=np.array(1, dtype)).sample(
            d, seed=42)
    radius = np.sqrt(2.)
    model_coefficients *= radius / tf.linalg.norm(model_coefficients)
    mask = tf.random_shuffle(tf.range(d)) < tf.to_int32(0.5 * tf.to_float(d))
    model_coefficients = tf.where(mask, model_coefficients,
                                  tf.zeros_like(model_coefficients))
    model_matrix = tfd.Normal(
        loc=np.array(0, dtype), scale=np.array(1, dtype)).sample(
            [n, d], seed=43)
    scale = tf.convert_to_tensor(scale, dtype)
    linear_response = tf.matmul(model_matrix,
                                model_coefficients[..., tf.newaxis])[..., 0]
    if link == 'linear':
      response = tfd.Normal(loc=linear_response, scale=scale).sample(seed=44)
    elif link == 'probit':
      response = tf.cast(
          tfd.Normal(loc=linear_response, scale=scale).sample(seed=44) > 0,
                     dtype)
    elif link == 'logit':
      response = tfd.Bernoulli(logits=linear_response).sample(seed=44)
    else:
      raise ValueError('unrecognized true link: {}'.format(link))
    return model_matrix, response, model_coefficients, mask

  with tf.Session() as sess:
    x_, y_, model_coefficients_true_, _ = sess.run(make_dataset(
        n=int(1e5), d=100, link='probit'))

    model = tfp.glm.Bernoulli()
    model_coefficients_start = tf.zeros(x_.shape[-1], np.float32)

    model_coefficients, is_converged, num_iter = tfp.glm.Func(
        model_matrix=tf.convert_to_tensor(x_),
        response=tf.convert_to_tensor(y_),
        model=model,
        model_coefficients_start=model_coefficients_start,
        l1_regularizer=800.,
        l2_regularizer=None,
        maximum_iterations=10,
        maximum_full_sweeps_per_iteration=10,
        tolerance=1e-6,
        learning_rate=None)

    model_coefficients_, is_converged_, num_iter_ = sess.run([
        model_coefficients, is_converged, num_iter])

    print("is_converged:", is_converged_)
    print("    num_iter:", num_iter_)
    print("\nLearned / True")
    print(np.concatenate(
        [[model_coefficients_], [model_coefficients_true_]], axis=0).T)

  # ==>
  # is_converged: True
  #     num_iter: 1
  #
  # Learned / True
  # [[ 0.          0.        ]
  #  [ 0.          0.        ]
  #  [ 0.          0.        ]
  #  [ 0.11195257  0.12484948]
  #  [ 0.          0.        ]
  #  [ 0.05191106  0.06394956]
  #  [-0.15090358 -0.15325639]
  #  [-0.18187316 -0.18825999]
  #  [-0.06140942 -0.07994166]
  #  [ 0.          0.        ]
  #  [ 0.          0.        ]
  #  [ 0.          0.        ]
  #  [ 0.14474444  0.15810856]
  #  [ 0.          0.        ]
  #  [-0.25249591 -0.24260855]
  #  [ 0.          0.        ]
  #  [ 0.          0.        ]
  #  [-0.03888761 -0.06755984]
  #  [ 0.          0.        ]
  #  [ 0.          0.        ]
  #  [ 0.          0.        ]
  #  [-0.0192222  -0.04169233]
  #  [ 0.          0.        ]
  #  [ 0.          0.        ]
  #  [ 0.01434913  0.03568212]
  #  [-0.11336883 -0.12873614]
  #  [ 0.          0.        ]
  #  [-0.24496339 -0.24048163]
  #  [ 0.          0.        ]
  #  [ 0.          0.        ]
  #  [ 0.04088281  0.06565224]
  #  [-0.12784363 -0.13359821]
  #  [ 0.05618424  0.07396613]
  #  [ 0.          0.        ]
  #  [ 0.          0.        ]
  #  [ 0.          0.        ]
  #  [ 0.         -0.01719233]
  #  [ 0.          0.        ]
  #  [ 0.          0.        ]
  #  [-0.00076072 -0.03607186]
  #  [ 0.21801499  0.21146794]
  #  [-0.02161094 -0.04031265]
  #  [ 0.0918689   0.10487888]
  #  [ 0.0106154   0.03233612]
  #  [-0.07817317 -0.09725142]
  #  [ 0.          0.        ]
  #  [ 0.          0.        ]
  #  [-0.23725343 -0.24194022]
  #  [ 0.          0.        ]
  #  [-0.08725718 -0.1048776 ]
  #  [ 0.          0.        ]
  #  [ 0.          0.        ]
  #  [-0.02114314 -0.04145789]
  #  [ 0.          0.        ]
  #  [ 0.          0.        ]
  #  [-0.02710908 -0.04590397]
  #  [ 0.15293184  0.15415154]
  #  [ 0.2114463   0.2088728 ]
  #  [-0.10969634 -0.12368613]
  #  [ 0.         -0.01505797]
  #  [-0.01140458 -0.03234904]
  #  [ 0.16051085  0.1680062 ]
  #  [ 0.09816848  0.11094204]
  ```

  #### References

  [1]: Jerome Friedman, Trevor Hastie and Rob Tibshirani. Regularization Paths
       for Generalized Linear Models via Coordinate Descent. _Journal of
       Statistical Software_, 33(1), 2010.
       https://www.jstatsoft.org/article/view/v033i01/v33i01.pdf

  [2]: Guo-Xun Yuan, Chia-Hua Ho and Chih-Jen Lin. An Improved GLMNET for
       L1-regularized Logistic Regression. _Journal of Machine Learning
       Research_, 13, 2012.
       http://www.jmlr.org/papers/volume13/yuan12a/yuan12a.pdf
  """
  arg_11 = [
      arg_0,
      arg_1,
      arg_3,
      arg_5,
      arg_6,
      arg_7,
      arg_8,
      # TODO(b/111925792): Replace `tolerance` arg with something like
      # `convergence_criteria_fn`.
      arg_4,
      arg_9,
  ]
  with tf.compat.v1.name_scope(arg_10, 'Func', arg_11):
    # TODO(b/111922388): Include dispersion and offset parameters.
    def _grad_neg_log_likelihood_and_fim_fn(arg_12):
      arg_13 = sparse_or_dense_matvecmul(arg_0, arg_12)
      arg_14, arg_15 = _grad_neg_log_likelihood_and_fim(
          arg_0, arg_13, arg_1, arg_2)
      return arg_14, arg_0, arg_15

    return tfp.optimizer.proximal_hessian_sparse_minimize(
        _grad_neg_log_likelihood_and_fim_fn,
        x_start=arg_3,
        arg_5=arg_5,
        arg_6=arg_6,
        arg_7=arg_7,
        arg_8=arg_8,
        arg_9=arg_9,
        arg_4=arg_4,
        arg_10=arg_10)


def Func(arg_0,
              arg_1,
              arg_2,
              arg_3=arg_4,
              arg_5=arg_6.float32):
  """Generate the mask for building an autoregressive dense layer."""
  # TODO(b/67594795): Better support of dynamic shape.
  arg_8 = np.zeros([arg_2, arg_1], arg_5=arg_5.as_numpy_dtype())
  arg_9 = _gen_slices(arg_0, arg_1, arg_2, arg_3=arg_3)
  for [arg_10, arg_11] in arg_9:
    arg_8[arg_10, arg_11] = 1
  return arg_8


def Func(arg_0,
                 arg_1,
                 arg_2=None,
                 arg_3=False,
                 arg_4=None,
                 arg_5=None,
                 arg_6=None,
                 *arg_7,  # pylint: disable=keyword-arg-before-vararg
                 **arg_8):
  """A autoregressively masked dense layer. Analogous to `tf.layers.dense`.

  See [Germain et al. (2015)][1] for detailed explanation.

  Arguments:
    inputs: Tensor input.
    units: Python `int` scalar representing the dimensionality of the output
      space.
    num_blocks: Python `int` scalar representing the number of blocks for the
      MADE masks.
    exclusive: Python `bool` scalar representing whether to zero the diagonal of
      the mask, used for the first layer of a MADE.
    kernel_initializer: Initializer function for the weight matrix.
      If `None` (default), weights are initialized using the
      `tf.glorot_random_initializer`.
    reuse: Python `bool` scalar representing whether to reuse the weights of a
      previous layer by the same name.
    name: Python `str` used to describe ops managed by this function.
    *args: `tf.layers.dense` arguments.
    **kwargs: `tf.layers.dense` keyword arguments.

  Returns:
    Output tensor.

  Raises:
    NotImplementedError: if rightmost dimension of `inputs` is unknown prior to
      graph execution.

  #### References

  [1]: Mathieu Germain, Karol Gregor, Iain Murray, and Hugo Larochelle. MADE:
       Masked Autoencoder for Distribution Estimation. In _International
       Conference on Machine Learning_, 2015. https://arxiv.org/abs/1502.03509
  """
  # TODO(b/67594795): Better support of dynamic shape.
  arg_9 = tf.compat.dimension_value(
      tensorshape_util.with_rank_at_least(arg_0.shape, 1)[-1])
  if arg_9 is None:
    raise NotImplementedError(
        "Rightmost dimension must be known prior to graph execution.")

  arg_10 = _gen_mask(arg_2, arg_9, arg_1,
                   MASK_EXCLUSIVE if arg_3 else MASK_INCLUSIVE).T

  if arg_4 is None:
    arg_4 = tf.compat.v1.glorot_normal_initializer()

  def masked_initializer(arg_11, arg_12=None, arg_13=None):
    return arg_10 * arg_4(arg_11, arg_12, arg_13)

  with tf.compat.v2.name_scope(arg_6 or "Func"):
    arg_14 = tf.compat.v1.layers.Dense(
        arg_1,
        arg_4=masked_initializer,
        kernel_constraint=lambda x: arg_10 * x,
        arg_6=arg_6,
        arg_12=dtype_util.base_dtype(arg_0.dtype),
        _scope=arg_6,
        _reuse=arg_5,
        *arg_7,  # pylint: disable=keyword-arg-before-vararg
        **arg_8)
    return arg_14.apply(arg_0)


def Func(arg_0, arg_1="left-to-right"):
  """Returns a degree vectors for the input."""
  if isinstance(arg_1, six.string_types):
    if arg_1 == "left-to-right":
      return np.arange(start=1, stop=arg_0 + 1)
    elif arg_1 == "right-to-left":
      return np.arange(start=arg_0, stop=0, step=-1)
    elif arg_1 == "random":
      arg_2 = np.arange(start=1, stop=arg_0 + 1)
      np.random.shuffle(arg_2)
      return arg_2
  elif np.all(np.sort(arg_1) == np.arange(1, arg_0 + 1)):
    return np.array(arg_1)

  raise ValueError("Invalid input order: '{}'.".format(arg_1))


def Func(arg_0):
  """Returns a list of binary mask matrices enforcing autoregressivity."""
  return [
      # Create input->hidden and hidden->hidden masks.
      arg_1[:, np.newaxis] <= arg_2
      for arg_1, arg_2 in zip(arg_0[:-1], arg_0[1:])
  ] + [
      # Create hidden->output mask.
      arg_0[-1][:, np.newaxis] < arg_0[0]
  ]


def Func(arg_0, arg_1):
  """Returns a masked version of the given initializer."""
  arg_1 = tf.keras.initializers.get(arg_1)
  def masked_initializer(arg_2, arg_3=None, arg_4=None):
    # If no `partition_info` is given, then don't pass it to `initializer`, as
    # `initializer` may be a `tf.compat.v2.initializers.Initializer` (which
    # don't accept a `partition_info` argument).
    if arg_4 is None:
      arg_5 = arg_1(arg_2, arg_3)
    else:
      arg_5 = arg_1(arg_2, arg_3, arg_4)
    return tf.cast(arg_0, arg_5.dtype) * arg_5
  return masked_initializer


def Func(arg_0, arg_1):
    """See tfkl.Layer.Func."""
    with tf.compat.v2.name_scope(arg_0.name or "AutoregressiveLayer_Func"):
      arg_1 = tf.convert_to_tensor(value=arg_1, dtype=arg_0.dtype, name="x")
      arg_2 = tf.shape(input=arg_1)
      # TODO(b/67594795): Better support for dynamic shapes.
      if tensorshape_util.rank(arg_1.shape) == 1:
        arg_1 = arg_1[tf.newaxis, ...]
      return tf.reshape(arg_0._network(arg_1),
                        tf.concat([arg_2, [arg_0._params]], axis=0))


def Func(arg_0, arg_1, arg_2, arg_3, arg_4, arg_5):
  """Sample a multinomial.

  The batch shape is given by broadcasting num_trials with
  remove_last_dimension(logits).

  Args:
    num_samples: Python int or singleton integer Tensor: number of multinomial
      samples to draw.
    num_classes: Python int or singleton integer Tensor: number of classes.
    logits: Floating Tensor with last dimension k, of (unnormalized) logit
      probabilities per class.
    num_trials: Tensor of number of categorical trials each multinomial consists
      of.  num_trials[..., tf.newaxis] must broadcast with logits.
    dtype: dtype at which to emit samples.
    seed: Random seed.

  Returns:
    samples: Tensor of given dtype and shape [n] + batch_shape + [k].
  """
  with tf.name_scope("multinomial.Func"):
    # broadcast the num_trials and logits to same shape
    arg_3 = tf.ones_like(
        arg_2[..., 0], arg_4=arg_3.dtype) * arg_3
    arg_2 = tf.ones_like(
        arg_3[..., tf.newaxis], arg_4=arg_2.dtype) * arg_2

    # flatten the total_count and logits
    # flat_logits has shape [B1B2...Bm, num_classes]
    arg_6 = tf.reshape(arg_2, [-1, arg_1])
    arg_7 = arg_0 * tf.reshape(arg_3, [-1])  # [B1B2...Bm]

    # Computes each logits and num_trials situation by map_fn.

    # Using just one batch tf.random.categorical call doesn't work because that
    # requires num_trials to be the same across all members of the batch of
    # logits.  This restriction makes sense for tf.random.categorical because
    # for it, num_trials is part of the returned shape.  However, the
    # multinomial sampler does not need that restriction, because it sums out
    # exactly that dimension.

    # One possibility would be to draw a batch categorical whose sample count is
    # max(num_trials) and mask out the excess ones.  However, if the elements of
    # num_trials vary widely, this can be wasteful of memory.

    # TODO(b/123763054, b/112152209): Revisit the possibility of writing this
    # with a batch categorical followed by batch unsorted_segment_sum, once both
    # of those work and are memory-efficient enough.
    def _sample_one_batch_member(arg_8):
      arg_2, arg_9 = arg_8[0], arg_8[1]  # [K], []
      # x has shape [1, num_cat_samples = num_samples * num_trials]
      arg_10 = tf.random.categorical(
          arg_2[tf.newaxis, ...], arg_9, arg_5=arg_5)
      arg_10 = tf.reshape(arg_10, shape=[arg_0, -1])  # [num_samples, num_trials]
      arg_10 = tf.one_hot(
          arg_10, depth=arg_1)  # [num_samples, num_trials, num_classes]
      arg_10 = tf.reduce_sum(input_tensor=arg_10, axis=-2)  # [num_samples, num_classes]
      return tf.cast(arg_10, arg_4=arg_4)

    arg_10 = tf.map_fn(
        _sample_one_batch_member, [arg_6, arg_7],
        arg_4=arg_4)  # [B1B2...Bm, num_samples, num_classes]

    # reshape the results to proper shape
    arg_10 = tf.transpose(a=arg_10, perm=[1, 0, 2])
    arg_11 = tf.concat([[arg_0],
                             tf.shape(input=arg_3), [arg_1]],
                            axis=0)
    arg_10 = tf.reshape(arg_10, arg_11)

    return arg_10


def Func(arg_0):
  """Build a zero-dimensional MVNDiag object."""
  arg_1 = tfd.MultivariateNormalDiag(
      scale_diag=tf.ones([0], arg_0=arg_0))
  arg_1.covariance = lambda: arg_1.variance()[..., tf.newaxis]
  return arg_1


def Func(arg_0):
  """Computes the number of edges on longest path from node to root."""
  def _explore(arg_1):
    if arg_1.depth < 0:
      arg_1.depth = ((1 + max([-1] + [_explore(arg_3[u])
                                  for u in arg_1.parents]))
                 if arg_1.parents else 0)
    return arg_1.depth
  arg_3 = {k: _Node(k, arg_1) for k, arg_1 in arg_0.items()}
  for arg_1 in arg_3.values():
    _explore(arg_1)
  return arg_3


def Func(arg_0):
  """Creates lists of callables suitable for JDSeq."""
  def _make(arg_1, arg_2):
    if arg_2 is None:
      return lambda *_: arg_1
    if not arg_2:
      return lambda *_: arg_1()
    def _fn(*arg_3):
      arg_4 = dict(zip(arg_2, reversed(arg_3[-len(arg_2):])))
      arg_4.pop('_', None)
      return arg_1(**arg_4)
    return _fn
  arg_0 = _convert_to_dict(arg_0)
  arg_5 = {k: (None if distribution_util.is_distribution_instance(v)
           else joint_distribution_sequential._get_required_args(v))  # pylint: disable=protected-access
       for k, v in arg_0.items()}
  arg_5 = _best_order(arg_5)
  arg_6, arg_7 = zip(*arg_5)
  arg_7 = tuple(None if a is None else tuple(a) for a in arg_7)
  arg_8 = tuple(_make(arg_0[name], parents)
                          for (name, parents) in arg_5)
  arg_1 = tuple(arg_0.get(n) for n in arg_6)
  return arg_1, arg_8, arg_7, arg_6


def Func(arg_0,
                       arg_1,
                       arg_2=None,
                       arg_3=1.,
                       arg_4='Func'):
    """Variational loss for the VGP.

    Given `observations` and `observation_index_points`, compute the
    negative variational lower bound as specified in [Hensman, 2013][1].

    Args:
      observations: `float` `Tensor` representing collection, or batch of
        collections, of observations corresponding to
        `observation_index_points`. Shape has the form `[b1, ..., bB, e]`, which
        must be brodcastable with the batch and example shapes of
        `observation_index_points`. The batch shape `[b1, ..., bB]` must be
        broadcastable with the shapes of all other batched parameters
        (`kernel.batch_shape`, `observation_index_points`, etc.).
      observation_index_points: `float` `Tensor` representing finite (batch of)
        vector(s) of points where observations are defined. Shape has the
        form `[b1, ..., bB, e1, f1, ..., fF]` where `F` is the number of feature
        dimensions and must equal `kernel.feature_ndims` and `e1` is the number
        (size) of index points in each batch (we denote it `e1` to distinguish
        it from the numer of inducing index points, denoted `e2` below). If
        set to `None` uses `index_points` as the origin for observations.
        Default value: None.
      kl_weight: Amount by which to scale the KL divergence loss between prior
        and posterior.
        Default value: 1.
      name: Python `str` name prefixed to Ops created by this class.
        Default value: "GaussianProcess".
    Returns:
      loss: Scalar tensor representing the negative variational lower bound.
        Can be directly used in a `tf.Optimizer`.
    Raises:
      ValueError: if `mean_fn` is not `None` and is not callable.

    #### References

    [1]: Hensman, J., Lawrence, N. "Gaussian Processes for Big Data", 2013
         https://arxiv.org/abs/1309.6835
    """

    with tf.name_scope(arg_4 or 'variational_gp_loss'):
      if arg_2 is None:
        arg_2 = arg_0._index_points
      arg_2 = tf.convert_to_tensor(
          value=arg_2, dtype=arg_0._dtype,
          arg_4='observation_index_points')
      arg_1 = tf.convert_to_tensor(
          value=arg_1, dtype=arg_0._dtype, arg_4='observations')
      arg_3 = tf.convert_to_tensor(
          value=arg_3, dtype=arg_0._dtype,
          arg_4='kl_weight')

      # The variational loss is a negative ELBO. The ELBO can be broken down
      # into three terms:
      #  1. a likelihood term
      #  2. a trace term arising from the covariance of the posterior predictive

      arg_5 = arg_0.kernel.matrix(arg_0._inducing_index_points,
                               arg_2)

      arg_6 = tf.linalg.LinearOperatorFullMatrix(arg_5)
      arg_7 = (arg_0._mean_fn(arg_2) +
             arg_6.matvec(arg_0._kzz_inv_varloc, adjoint=True))

      arg_8 = independent.Independent(
          normal.Normal(
              arg_7=arg_7,
              scale=tf.sqrt(arg_0._observation_noise_variance + arg_0._jitter),
              arg_4='NormalLikelihood'),
          reinterpreted_batch_ndims=1)
      arg_9 = arg_8.log_prob(arg_1)

      arg_10 = tf.linalg.LinearOperatorLowerTriangular(arg_0._chol_kzz)
      arg_11 = arg_10.solve(arg_5)
      arg_12 = arg_10.solve(arg_11, adjoint=True)

      arg_13 = tf.linalg.diag_part(
          arg_0.kernel.matrix(
              arg_2, arg_2))
      arg_14 = (
          tf.reduce_sum(input_tensor=arg_13, axis=-1) -
          tf.reduce_sum(input_tensor=arg_11 ** 2, axis=[-2, -1]))

      # Tr(SB)
      # where S = A A.T, A = variational_inducing_observations_scale
      # and B = Kzz^-1 Kzx Kzx.T Kzz^-1
      #
      # Now Tr(SB) = Tr(A A.T Kzz^-1 Kzx Kzx.T Kzz^-1)
      #            = Tr(A.T Kzz^-1 Kzx Kzx.T Kzz^-1 A)
      #            = sum_ij (A.T Kzz^-1 Kzx)_{ij}^2
      arg_15 = tf.reduce_sum(
          input_tensor=(
              arg_0._variational_inducing_observations_posterior.scale.matmul(
                  arg_12) ** 2),
          axis=[-2, -1])

      arg_16 = (.5 * (arg_14 + arg_15) /
                    arg_0._observation_noise_variance)

      arg_17 = gaussian_process.GaussianProcess(
          kernel=arg_0._kernel,
          mean_fn=arg_0._mean_fn,
          index_points=arg_0._inducing_index_points,
          observation_noise_variance=arg_0._observation_noise_variance)

      arg_18 = arg_3 * kullback_leibler.kl_divergence(
          arg_0._variational_inducing_observations_posterior,
          arg_17)

      arg_19 = (arg_9 - arg_16 - arg_18)

      return -tf.reduce_mean(input_tensor=arg_19)


def Func(
      arg_0,
      arg_1,
      arg_2,
      arg_3,
      arg_4,
      arg_5=None,
      arg_6=1e-6,
      arg_7=None):
    """Model selection for optimal variational hyperparameters.

    Given the full training set (parameterized by `observations` and
    `observation_index_points`), compute the optimal variational
    location and scale for the VGP. This is based of the method suggested
    in [Titsias, 2009][1].

    Args:
      kernel: `PositiveSemidefiniteKernel`-like instance representing the
        GP's covariance function.
      inducing_index_points: `float` `Tensor` of locations of inducing points in
        the index set. Shape has the form `[b1, ..., bB, e2, f1, ..., fF]`, just
        like `observation_index_points`. The batch shape components needn't be
        identical to those of `observation_index_points`, but must be broadcast
        compatible with them.
      observation_index_points: `float` `Tensor` representing finite (batch of)
        vector(s) of points where observations are defined. Shape has the
        form `[b1, ..., bB, e1, f1, ..., fF]` where `F` is the number of feature
        dimensions and must equal `kernel.feature_ndims` and `e1` is the number
        (size) of index points in each batch (we denote it `e1` to distinguish
        it from the numer of inducing index points, denoted `e2` below).
      observations: `float` `Tensor` representing collection, or batch of
        collections, of observations corresponding to
        `observation_index_points`. Shape has the form `[b1, ..., bB, e]`, which
        must be brodcastable with the batch and example shapes of
        `observation_index_points`. The batch shape `[b1, ..., bB]` must be
        broadcastable with the shapes of all other batched parameters
        (`kernel.batch_shape`, `observation_index_points`, etc.).
      observation_noise_variance: `float` `Tensor` representing the variance
        of the noise in the Normal likelihood distribution of the model. May be
        batched, in which case the batch shape must be broadcastable with the
        shapes of all other batched parameters (`kernel.batch_shape`,
        `index_points`, etc.).
        Default value: `0.`
      mean_fn: Python `callable` that acts on index points to produce a (batch
        of) vector(s) of mean values at those index points. Takes a `Tensor` of
        shape `[b1, ..., bB, f1, ..., fF]` and returns a `Tensor` whose shape is
        (broadcastable with) `[b1, ..., bB]`. Default value: `None` implies
        constant zero function.
      jitter: `float` scalar `Tensor` added to the diagonal of the covariance
        matrix to ensure positive definiteness of the covariance matrix.
        Default value: `1e-6`.
      name: Python `str` name prefixed to Ops created by this class.
        Default value: "Func".
    Returns:
      loc, scale: Tuple representing the variational location and scale.
    Raises:
      ValueError: if `mean_fn` is not `None` and is not callable.

    #### References

    [1]: Titsias, M. "Variational Model Selection for Sparse Gaussian Process
         Regression", 2009.
         http://proceedings.mlr.press/v5/titsias09a/titsias09a.pdf
    """

    with tf.name_scope(arg_7 or 'Func'):
      arg_8 = dtype_util.common_dtype(
          [arg_1,
           arg_2,
           arg_3,
           arg_4,
           arg_6], tf.float32)

      arg_1 = tf.convert_to_tensor(
          value=arg_1,
          arg_8=arg_8, arg_7='inducing_index_points')
      arg_2 = tf.convert_to_tensor(
          value=arg_2, arg_8=arg_8,
          arg_7='observation_index_points')
      arg_3 = tf.convert_to_tensor(
          value=arg_3, arg_8=arg_8, arg_7='observations')
      arg_4 = tf.convert_to_tensor(
          value=arg_4,
          arg_8=arg_8,
          arg_7='observation_noise_variance')
      arg_6 = tf.convert_to_tensor(
          value=arg_6, arg_8=arg_8, arg_7='jitter')

      # Default to a constant zero function.
      if arg_5 is None:
        arg_5 = lambda x: tf.zeros([1], arg_8=arg_8)
      else:
        if not callable(arg_5):
          raise ValueError('`mean_fn` must be a Python callable')

      # z are the inducing points and x are the observation index points.
      arg_9 = arg_0.matrix(arg_1, arg_1)
      arg_10 = arg_0.matrix(arg_1, arg_2)

      arg_11 = tf.math.reciprocal(arg_4)

      arg_12 = _add_diagonal_shift(
          arg_9 + arg_11 * tf.matmul(arg_10, arg_10, adjoint_b=True),
          arg_6)

      arg_13 = tf.linalg.cholesky(arg_12)

      arg_14 = tf.linalg.LinearOperatorFullMatrix(arg_10)
      arg_15 = arg_14.matvec(
          arg_3 - arg_5(arg_2))
      arg_16 = tf.linalg.LinearOperatorFullMatrix(arg_9)
      arg_17 = (arg_5(arg_1) +
             arg_11 * arg_16.matvec(
                 _solve_cholesky_factored_system_vec(arg_13, arg_15)))

      arg_18 = tf.linalg.LinearOperatorLowerTriangular(
          arg_13)
      arg_19 = arg_18.solve(arg_9)

      return arg_17, arg_19


def Func(arg_0):
  """Build utility method to compute whether the season is changing."""
  arg_1 = np.sum(arg_0)
  arg_2 = np.cumsum(np.ravel(arg_0)) - 1
  def is_last_day_of_season(arg_3):
    arg_4 = dist_util.maybe_get_static_value(arg_3)
    if arg_4 is not None:  # static case
      arg_5 = arg_4 % arg_1
      return any(arg_5 == arg_2)
    else:
      arg_5 = tf.math.floormod(arg_3, arg_1)
      return tf.reduce_any(
          input_tensor=tf.equal(arg_5, arg_2))
  return is_last_day_of_season


def Func(
    arg_0, arg_1, arg_2,
    arg_3=None, arg_4=None):
  """Build a function computing transitions for a seasonal effect model."""

  with tf.compat.v1.name_scope('Func'):
    # If the season is changing, the transition matrix permutes the latent
    # state to shift all seasons up by a dimension, and sends the current
    # season's effect to the bottom.
    arg_5 = np.concatenate(
        [np.arange(1, arg_0), [0]], axis=0)
    arg_6 = tf.constant(
        np.eye(arg_0)[arg_5], arg_2=arg_2)

    # Optionally transform the transition matrix into a reparameterized space,
    # enforcing the zero-sum constraint for ConstrainedSeasonalStateSpaceModel.
    if arg_3 is not None:
      arg_6 = tf.matmul(
          arg_3,
          tf.matmul(arg_6, arg_4))

    arg_7 = tf.eye(
        tf.shape(input=arg_6)[-1], arg_2=arg_2)

    def seasonal_transition_matrix(arg_8):
      return tf.linalg.LinearOperatorFullMatrix(
          matrix=dist_util.pick_scalar_condition(
              arg_1(arg_8),
              arg_6,
              arg_7))

  return seasonal_transition_matrix


def Func(
    arg_0, arg_1, arg_2):
  """Build the transition noise model for a SeasonalStateSpaceModel."""

  # If the current season has just ended, increase the variance of its effect
  # following drift_scale. (the just-ended seasonal effect will always be the
  # bottom element of the vector). Otherwise, do nothing.
  arg_3 = tf.stack(
      [tf.zeros_like(arg_0)] * (arg_1 - 1) + [arg_0],
      axis=-1)
  def seasonal_transition_noise(arg_4):
    arg_5 = dist_util.pick_scalar_condition(
        arg_2(arg_4),
        arg_3,
        tf.zeros_like(arg_3))
    return tfd.MultivariateNormalDiag(
        loc=tf.zeros(arg_1, dtype=arg_0.dtype),
        scale_diag=arg_5)
  return seasonal_transition_noise


def Func(
    arg_0, arg_1, arg_2):
  """Returns `True` if given observation data is empty.

  Emptiness means either
    1. Both `observation_index_points` and `observations` are `None`, or
    2. the "number of observations" shape is 0. The shape of
    `observation_index_points` is `[..., N, f1, ..., fF]`, where `N` is the
    number of observations and the `f`s are feature dims. Thus, we look at the
    shape element just to the left of the leftmost feature dim. If that shape is
    zero, we consider the data empty.

  We don't check the shape of observations; validations are checked elsewhere in
  the calling code, to ensure these shapes are consistent.

  Args:
    feature_ndims: the number of feature dims, as reported by the GP kernel.
    observation_index_points: the observation data locations in the index set.
    observations: the observation data.

  Returns:
    is_empty: True if the data were deemed to be empty.
  """
  # If both input locations and observations are `None`, we consider this
  # "empty" observation data.
  if arg_1 is None and arg_2 is None:
    return True
  arg_3 = tf.compat.dimension_value(
      arg_1.shape[-(arg_0 + 1)])
  if arg_3 is not None and arg_3 == 0:
    return True
  return False


def Func(
    arg_0, arg_1, arg_2):
  """Ensure that observation data and locations have consistent shapes.

  This basically means that the batch shapes are broadcastable. We can only
  ensure this when those shapes are fully statically defined.

  Args:
    kernel: The GP kernel.
    observation_index_points: the observation data locations in the index set.
    observations: the observation data.

  Raises:
    ValueError: if the observations' batch shapes are not broadcastable.
  """
  # Check that observation index points and observation counts broadcast.
  arg_3 = arg_0.feature_ndims
  if (tensorshape_util.is_fully_defined(
      arg_1.shape[:-arg_3]) and
      tensorshape_util.is_fully_defined(arg_2.shape)):
    arg_4 = arg_1.shape[:-arg_3]
    arg_5 = arg_2.shape
    try:
      tf.broadcast_static_shape(arg_4, arg_5)
    except ValueError:
      # Re-raise with our own more contextual error message.
      raise ValueError(
          'Observation index point and observation counts are not '
          'broadcastable: {} and {}, respectively.'.format(
              arg_4, arg_5))


def Func(arg_0, arg_1, arg_2, arg_3="float"):
        """
        Add a learning rate scheduler to the contained `schedules`

        :param scheduler: learning rate scheduler to be Func
        :param max_iteration: iteration numbers this scheduler will run
        """
        return callBigDlFunc(arg_3, "FuncScheduler", arg_0.value, arg_1, arg_2)


def Func(arg_0, arg_1,
                       arg_2, arg_3=True):
        """
        Configure checkpoint settings.

        :param checkpoint_trigger: the interval to write snapshots
        :param checkpoint_path: the path to write snapshots into
        :param isOverWrite: whether to overwrite existing snapshots in path.default is True
        """
        if not os.path.exists(arg_2):
            mkpath(arg_2)
        callBigDlFunc(arg_0.bigdl_type, "setCheckPoint", arg_0.value,
                      arg_1, arg_2, arg_3)


def Func(arg_0, arg_1, arg_2):
        """
        Configure constant clipping settings.

        :param min_value: the minimum value to clip by
        :param max_value: the maxmimum value to clip by
        """
        callBigDlFunc(arg_0.bigdl_type, "setConstantClip", arg_0.value, arg_1, arg_2)


def Func(arg_0):
        """
        Do an optimization.
        """
        arg_1 = callJavaFunc(arg_0.value.Func)
        from bigdl.nn.layer import Layer
        return Layer.of(arg_1)


def Func(arg_0, arg_1):
        """
        Set validation summary. A ValidationSummary object contains information
        necessary for the optimizer to know how often the logs are recorded,
        where to store the logs and how to retrieve them, etc. For details,
        refer to the docs of ValidationSummary.

        :param summary: a ValidationSummary object

        """
        callBigDlFunc(arg_0.bigdl_type, "setValSummary", arg_0.value,
                      arg_1)
        return arg_0


def Func(arg_0="./data/news20/"):
    """
    Parse or download news20 if source_dir is empty.

    :param source_dir: The directory storing news data.
    :return: A list of (tokens, label)
    """
    arg_1 = download_news20(arg_0)
    arg_2 = []  # list of text samples
    arg_3 = 0
    for arg_4 in sorted(os.listdir(arg_1)):
        arg_5 = os.path.join(arg_1, arg_4)
        arg_3 += 1
        if os.path.isdir(arg_5):
            for arg_6 in sorted(os.listdir(arg_5)):
                if arg_6.isdigit():
                    arg_7 = os.path.join(arg_5, arg_6)
                    if sys.version_info < (3,):
                        arg_8 = open(arg_7)
                    else:
                        arg_8 = open(arg_7, encoding='latin-1')
                    arg_9 = arg_8.read()
                    arg_2.append((arg_9, arg_3))
                    arg_8.close()

    print('Found %s texts.' % len(arg_2))
    return arg_2


def Func(arg_0="./data/news20/", arg_1=100):
    """
    Parse or download the pre-trained glove word2vec if source_dir is empty.

    :param source_dir: The directory storing the pre-trained word2vec
    :param dim: The dimension of a vector
    :return: A dict mapping from word to vector
    """
    arg_2 = download_glove_w2v(arg_0)
    arg_3 = os.path.join(arg_2, "glove.6B.%sd.txt" % arg_1)
    if sys.version_info < (3,):
        arg_4 = open(arg_3)
    else:
        arg_4 = open(arg_3, encoding='latin-1')
    arg_5 = {}
    for arg_6 in arg_4.readlines():
        arg_7 = arg_6.split(" ")
        arg_5[arg_7[0]] = [float(i) for i in arg_7[1:]]
    arg_4.close()
    return arg_5


def Func(arg_0, arg_1, arg_2=None, arg_3=32, arg_4=10, arg_5=None, arg_6=True):
        """
        Train a model for a fixed number of epochs on a dataset.

        # Arguments
        x: Input data. A Numpy array or RDD of Sample or Image DataSet.
        y: Labels. A Numpy array. Default is None if x is already RDD of Sample or Image DataSet.
        batch_size: Number of samples per gradient update.
        nb_epoch: Number of iterations to train.
        validation_data: Tuple (x_val, y_val) where x_val and y_val are both Numpy arrays.
                         Or RDD of Sample. Default is None if no validation is involved.
        distributed: Boolean. Whether to train the model in distributed mode or local mode.
                     Default is True. In local mode, x and y must both be Numpy arrays.
        """
        if arg_6:
            if isinstance(arg_1, np.ndarray) and isinstance(arg_2, np.ndarray):
                arg_7 = to_sample_rdd(arg_1, arg_2)
                if arg_5:
                    arg_5 = to_sample_rdd(*arg_5)
            elif (isinstance(arg_1, RDD) or isinstance(arg_1, DataSet)) and not arg_2:
                arg_7 = arg_1
            else:
                raise TypeError("Unsupported training data type: %s" % type(arg_1))
            callBigDlFunc(arg_0.bigdl_type, "Func",
                          arg_0.value,
                          arg_7,
                          arg_3,
                          arg_4,
                          arg_5)
        else:
            if arg_5:
                arg_8 = [JTensor.from_ndarray(arg_1) for arg_1 in to_list(arg_5[0])]
                arg_9 = JTensor.from_ndarray(arg_5[1])
            else:
                arg_8, arg_9 = None, None
            callBigDlFunc(arg_0.bigdl_type, "Func",
                          arg_0.value,
                          [JTensor.from_ndarray(arg_1) for arg_1 in to_list(arg_1)],
                          JTensor.from_ndarray(arg_2),
                          arg_3,
                          arg_4,
                          arg_8,
                          arg_9,
                          multiprocessing.cpu_count())


def Func(arg_0, arg_1, arg_2=None, arg_3=32):
        """
        Evaluate a model on a given dataset in distributed mode.

        # Arguments
        x: Input data. A Numpy array or RDD of Sample.
        y: Labels. A Numpy array. Default is None if x is already RDD of Sample.
        batch_size: Number of samples per gradient update.
        """
        if isinstance(arg_1, np.ndarray) and isinstance(arg_2, np.ndarray):
            arg_4 = to_sample_rdd(arg_1, arg_2)
        elif isinstance(arg_1, RDD) and not arg_2:
            arg_4 = arg_1
        else:
            raise TypeError("Unsupported evaluation data type: %s" % type(arg_1))
        return callBigDlFunc(arg_0.bigdl_type, "Func",
                             arg_0.value,
                             arg_4,
                             arg_3)


def Func(arg_0, arg_1="train", arg_2="/tmp/mnist"):
    """
    Get mnist dataset and parallelize into RDDs.
    Data would be downloaded automatically if it doesn't present at the specific location.

    :param sc: SparkContext.
    :param data_type: "train" for training data and "test" for testing data.
    :param location: Location to store mnist dataset.
    :return: RDD of (features: ndarray, label: ndarray).
    """
    (arg_3, arg_4) = mnist.read_data_sets(arg_2, arg_1)
    arg_3 = arg_0.parallelize(arg_3)
    arg_4 = arg_0.parallelize(arg_4 + 1)  # Target start from 1 in BigDL
    arg_5 = arg_3.zip(arg_4)
    return arg_5