def Func(arg_0=None,
                    arg_1=None,
                    arg_2=None,
                    arg_3=None,
                    arg_4=False,
                    arg_5=False,
                    arg_6=None,
                    arg_7=None):
  """Creates a LinearOperator representing a diagonal matrix.

  Args:
    loc: Floating-point `Tensor`. This is used for inferring shape in the case
      where only `scale_identity_multiplier` is set.
    scale_diag: Floating-point `Tensor` representing the diagonal matrix.
      `scale_diag` has shape [N1, N2, ...  k], which represents a k x k diagonal
      matrix. When `None` no diagonal term is added to the LinearOperator.
    scale_identity_multiplier: floating point rank 0 `Tensor` representing a
      scaling done to the identity matrix. When `scale_identity_multiplier =
      scale_diag = scale_tril = None` then `scale += IdentityMatrix`. Otherwise
      no scaled-identity-matrix is added to `scale`.
    shape_hint: scalar integer `Tensor` representing a hint at the dimension of
      the identity matrix when only `scale_identity_multiplier` is set.
    validate_args: Python `bool` indicating whether arguments should be checked
      for correctness.
    assert_positive: Python `bool` indicating whether LinearOperator should be
      checked for being positive definite.
    name: Python `str` name given to ops managed by this object.
    dtype: TF `DType` to prefer when converting args to `Tensor`s. Else, we fall
      back to a compatible dtype across all of `loc`, `scale_diag`, and
      `scale_identity_multiplier`.

  Returns:
    `LinearOperator` representing a lower triangular matrix.

  Raises:
    ValueError:  If only `scale_identity_multiplier` is set and `loc` and
      `shape_hint` are both None.
  """

  def _maybe_attach_assertion(arg_8):
    if not arg_4:
      return arg_8
    if arg_5:
      return with_dependencies([
          assert_util.assert_positive(
              arg_8, message="diagonal part must be positive"),
      ], arg_8)
    return with_dependencies([
        assert_util.assert_none_equal(
            arg_8, tf.zeros([], arg_8.dtype), message="diagonal part must be non-zero")
    ], arg_8)

  with tf.name_scope(arg_6 or "Func"):
    if arg_7 is None:
      arg_7 = dtype_util.common_dtype(
          [arg_0, arg_1, arg_2],
          preferred_dtype=tf.float32)
    arg_0 = _convert_to_tensor(arg_0, arg_6="loc", arg_7=arg_7)
    arg_1 = _convert_to_tensor(arg_1, arg_6="scale_diag", arg_7=arg_7)
    arg_2 = _convert_to_tensor(
        arg_2,
        arg_6="scale_identity_multiplier",
        arg_7=arg_7)

    if arg_1 is not None:
      if arg_2 is not None:
        arg_1 += arg_2[..., tf.newaxis]
      return tf.linalg.LinearOperatorDiag(
          diag=_maybe_attach_assertion(arg_1),
          is_non_singular=True,
          is_self_adjoint=True,
          is_positive_definite=arg_5)

    if arg_0 is None and arg_3 is None:
      raise ValueError("Cannot infer `event_shape` unless `loc` or "
                       "`shape_hint` is specified.")

    arg_9 = arg_3
    del arg_3
    if arg_9 is None:
      arg_9 = tf.compat.dimension_value(arg_0.shape[-1])
      if arg_9 is None:
        arg_9 = tf.shape(input=arg_0)[-1]

    if arg_2 is None:
      return tf.linalg.LinearOperatorIdentity(
          arg_9=arg_9,
          arg_7=arg_7,
          is_self_adjoint=True,
          is_positive_definite=True,
          assert_proper_shapes=arg_4)

    return tf.linalg.LinearOperatorScaledIdentity(
        arg_9=arg_9,
        multiplier=_maybe_attach_assertion(arg_2),
        is_non_singular=True,
        is_self_adjoint=True,
        is_positive_definite=arg_5,
        assert_proper_shapes=arg_4)


def Func(arg_0, arg_1, arg_2="Func"):
  """Infer distribution batch and event shapes from a location and scale.

  Location and scale family distributions determine their batch/event shape by
  broadcasting the `loc` and `scale` args.  This helper does that broadcast,
  statically if possible.

  Batch shape broadcasts as per the normal rules.
  We allow the `loc` event shape to broadcast up to that of `scale`.  We do not
  allow `scale`'s event shape to change.  Therefore, the last dimension of `loc`
  must either be size `1`, or the same as `scale.range_dimension`.

  See `MultivariateNormalLinearOperator` for a usage example.

  Args:
    loc: `Tensor` (already converted to tensor) or `None`. If `None`, or
      `rank(loc)==0`, both batch and event shape are determined by `scale`.
    scale:  A `LinearOperator` instance.
    name:  A string name to prepend to created ops.

  Returns:
    batch_shape:  `TensorShape` (if broadcast is done statically), or `Tensor`.
    event_shape:  `TensorShape` (if broadcast is done statically), or `Tensor`.

  Raises:
    ValueError:  If the last dimension of `loc` is determined statically to be
      different than the range of `scale`.
  """
  if arg_0 is not None and tensorshape_util.rank(arg_0.shape) == 0:
    arg_0 = None  # scalar loc is irrelevant to determining batch/event shape.
  with tf.name_scope(arg_2):
    # Get event shape.
    arg_3 = arg_1.range_dimension_tensor()
    arg_4 = tf.get_static_value(arg_3)
    arg_5 = (None if arg_0 is None
                       else tf.compat.dimension_value(arg_0.shape[-1]))

    if arg_4 is not None and arg_5 is not None:
      # Static check that event shapes match.
      if arg_5 != 1 and arg_5 != arg_4:
        raise ValueError(
            "Event size of 'scale' ({}) could not be broadcast up to that "
            "of 'loc' ({}).".format(arg_4, arg_5))
    elif arg_5 is not None and arg_5 != 1:
      arg_4 = arg_5

    if arg_4 is None:
      arg_6 = arg_3[tf.newaxis]
    else:
      arg_6 = tf.convert_to_tensor(
          value=np.reshape(arg_4, [1]),
          dtype=tf.int32,
          arg_2="event_shape")

    # Get batch shape.
    arg_7 = arg_1.batch_shape_tensor()
    if arg_0 is not None:
      arg_8 = tensorshape_util.with_rank_at_least(arg_0.shape, 1)[:-1]
      if tensorshape_util.rank(
          arg_0.shape) is None or not tensorshape_util.is_fully_defined(
              arg_8):
        arg_8 = tf.shape(input=arg_0)[:-1]
      else:
        arg_8 = tf.convert_to_tensor(
            value=arg_8, dtype=tf.int32, arg_2="loc_batch_shape")
      # This is defined in the core util module.
      arg_7 = prefer_static_broadcast_shape(arg_7, arg_8)  # pylint: disable=undefined-variable
      arg_7 = tf.convert_to_tensor(
          value=arg_7, dtype=tf.int32, arg_2="batch_shape")

    return arg_7, arg_6


def Func(arg_0):
  """Returns `True` if `scale` is a `LinearOperator` that is known to be diag.

  Args:
    scale:  `LinearOperator` instance.

  Returns:
    Python `bool`.

  Raises:
    TypeError:  If `scale` is not a `LinearOperator`.
  """
  if not isinstance(arg_0, tf.linalg.LinearOperator):
    raise TypeError("Expected argument 'scale' to be instance of LinearOperator"
                    ". Found: %s" % arg_0)
  return (isinstance(arg_0, tf.linalg.LinearOperatorIdentity) or
          isinstance(arg_0, tf.linalg.LinearOperatorScaledIdentity) or
          isinstance(arg_0, tf.linalg.LinearOperatorDiag))


def Func(arg_0, arg_1, arg_2, arg_3=None):
  """Convenience function that chooses one of two values based on the predicate.

  This utility is equivalent to a version of `tf.where` that accepts only a
  scalar predicate and computes its result statically when possible. It may also
  be used in place of `tf.cond` when both branches yield a `Tensor` of the same
  shape; the operational difference is that `tf.cond` uses control flow to
  evaluate only the branch that's needed, while `tf.where` (and thus
  this method) may evaluate both branches before the predicate's truth is known.
  This means that `tf.cond` is preferred when one of the branches is expensive
  to evaluate (like performing a large matmul), while this method is preferred
  when both branches are cheap, e.g., constants. In the latter case, we expect
  this method to be substantially faster than `tf.cond` on GPU and to give
  similar performance on CPU.

  Args:
    pred: Scalar `bool` `Tensor` predicate.
    true_value: `Tensor` to return if `pred` is `True`.
    false_value: `Tensor` to return if `pred` is `False`. Must have the same
      shape as `true_value`.
    name: Python `str` name given to ops managed by this object.

  Returns:
    result: a `Tensor` (or `Tensor`-convertible Python value) equal to
      `true_value` if `pred` evaluates to `True` and `false_value` otherwise.
      If the condition can be evaluated statically, the result returned is one
      of the input Python values, with no graph side effects.
  """
  with tf.name_scope(arg_3 or "Func"):
    arg_0 = tf.convert_to_tensor(
        value=arg_0, dtype_hint=tf.bool, arg_3="pred")
    arg_1 = tf.convert_to_tensor(value=arg_1, arg_3="true_value")
    arg_2 = tf.convert_to_tensor(value=arg_2, arg_3="false_value")
    arg_4 = tf.get_static_value(arg_0)
    if arg_4 is None:
      return tf.where(arg_0, arg_1, arg_2)
    return arg_1 if arg_4 else arg_2


def Func(arg_0, arg_1, arg_2):
  """Move a single tensor dimension within its shape.

  This is a special case of `tf.transpose()`, which applies
  arbitrary permutations to tensor dimensions.

  Args:
    x: Tensor of rank `ndims`.
    source_idx: Integer index into `x.shape` (negative indexing is supported).
    dest_idx: Integer index into `x.shape` (negative indexing is supported).

  Returns:
    x_perm: Tensor of rank `ndims`, in which the dimension at original
     index `source_idx` has been moved to new index `dest_idx`, with
     all other dimensions retained in their original order.

  Example:

  ```python
  x = tf.placeholder(shape=[200, 30, 4, 1, 6])
  x_perm = _Func(x, 1, 1) # no-op
  x_perm = _Func(x, 0, 3) # result shape [30, 4, 1, 200, 6]
  x_perm = _Func(x, 0, -2) # equivalent to previous
  x_perm = _Func(x, 4, 2) # result shape [200, 30, 6, 4, 1]
  ```
  """
  arg_3 = prefer_static_rank(arg_0)
  arg_4 = dtype_util.common_dtype([arg_1, arg_2],
                                  preferred_dtype=tf.int32)
  arg_1 = tf.convert_to_tensor(value=arg_1, arg_4=arg_4)
  arg_2 = tf.convert_to_tensor(value=arg_2, arg_4=arg_4)

  # Handle negative indexing.
  arg_1 = pick_scalar_condition(arg_1 < 0, arg_3 + arg_1,
                                     arg_1)
  arg_2 = pick_scalar_condition(arg_2 < 0, arg_3 + arg_2, arg_2)

  # Construct the appropriate permutation of dimensions, depending
  # whether the source is before or after the destination.
  def move_left_permutation():
    return prefer_static_value(
        tf.concat([
            tf.range(0, arg_2, arg_4=arg_4), [arg_1],
            tf.range(arg_2, arg_1, arg_4=arg_4),
            tf.range(arg_1 + 1, arg_3, arg_4=arg_4)
        ],
                  axis=0))

  def move_right_permutation():
    return prefer_static_value(
        tf.concat([
            tf.range(0, arg_1, arg_4=arg_4),
            tf.range(arg_1 + 1, arg_2 + 1, arg_4=arg_4), [arg_1],
            tf.range(arg_2 + 1, arg_3, arg_4=arg_4)
        ],
                  axis=0))

  def x_permuted():
    return tf.transpose(
        a=arg_0,
        perm=prefer_static.cond(arg_1 < arg_2,
                                move_right_permutation,
                                move_left_permutation))

  # One final conditional to handle the special case where source
  # and destination indices are equal.
  return prefer_static.cond(tf.equal(arg_1, arg_2),
                            lambda: arg_0, x_permuted)


def Func(
    arg_0, arg_1="Func"):
  """Assert x is a non-negative tensor, and optionally of integers."""
  with tf.name_scope(arg_1):
    arg_0 = tf.convert_to_tensor(value=arg_0, arg_1="x")
    arg_2 = [
        assert_util.assert_non_negative(
            arg_0, message="'{}' must be non-negative.".format(arg_0)),
    ]
    if not dtype_util.is_integer(arg_0.dtype):
      arg_2 += [
          assert_integer_form(
              arg_0,
              message="'{}' cannot contain fractional components.".format(arg_0)),
      ]
    return with_dependencies(arg_2, arg_0)


def Func(arg_0):
  """Helper returning True if dtype is known to be unsigned."""
  return {
      tf.bool: True,
      tf.uint8: True,
      tf.uint16: True,
  }.get(arg_0.base_dtype, False)


def Func(arg_0):
  """Helper returning True if dtype is known to be signed."""
  return {
      tf.float16: True,
      tf.float32: True,
      tf.float64: True,
      tf.int8: True,
      tf.int16: True,
      tf.int32: True,
      tf.int64: True,
  }.get(arg_0.base_dtype, False)


def Func(arg_0):
  """Helper returning the largest integer exactly representable by dtype."""
  if not _is_known_dtype(arg_0):
    raise TypeError("Unrecognized dtype: {}".format(arg_0.name))
  if arg_0.is_floating:
    return int(2**(np.finfo(arg_0.as_numpy_dtype).nmant + 1))
  if arg_0.is_integer:
    return np.iinfo(arg_0.as_numpy_dtype).max
  if arg_0.base_dtype == tf.bool:
    return int(1)
  # We actually can't land here but keep the case for completeness.
  raise TypeError("Unrecognized dtype: {}".format(arg_0.name))


def Func(arg_0):
  """Helper returning the smallest integer exactly representable by dtype."""
  if not _is_known_dtype(arg_0):
    raise TypeError("Unrecognized dtype: {}".format(arg_0.name))
  if _is_known_unsigned_by_dtype(arg_0):
    return 0
  return -1 * _largest_integer_by_dtype(arg_0)


def Func(arg_0, arg_1, arg_2="Func"):
  """Circularly moves dims left or right.

  Effectively identical to:

  ```python
  numpy.transpose(x, numpy.roll(numpy.arange(len(x.shape)), shift))
  ```

  When `validate_args=False` additional graph-runtime checks are
  performed. These checks entail moving data from to GPU to CPU.

  Example:

  ```python
  x = tf.random_normal([1, 2, 3, 4])  # Tensor of shape [1, 2, 3, 4].
  Func(x, -1).shape == [2, 3, 4, 1]
  Func(x, -2).shape == [3, 4, 1, 2]
  Func(x,  1).shape == [4, 1, 2, 3]
  Func(x,  2).shape == [3, 4, 1, 2]
  Func(x,  7).shape == Func(x, 3).shape  # [2, 3, 4, 1]
  Func(x, -7).shape == Func(x, -3).shape  # [4, 1, 2, 3]
  ```

  Args:
    x: `Tensor`.
    shift: `Tensor`. Number of dimensions to transpose left (shift<0) or
      transpose right (shift>0).
    name: Python `str`. The name to give this op.

  Returns:
    rotated_x: Input `Tensor` with dimensions circularly rotated by shift.

  Raises:
    TypeError: if shift is not integer type.
  """
  with tf.name_scope(arg_2):
    arg_0 = tf.convert_to_tensor(value=arg_0, arg_2="x")
    arg_1 = tf.convert_to_tensor(value=arg_1, arg_2="shift")
    # We do not assign back to preserve constant-ness.
    assert_util.assert_integer(arg_1)
    arg_3 = tf.get_static_value(arg_1)
    arg_4 = tensorshape_util.rank(arg_0.shape)
    if arg_4 is not None and arg_3 is not None:
      if arg_4 < 2:
        return arg_0
      arg_3 = np.sign(arg_3) * (
          abs(arg_3) % arg_4)
      if arg_3 == 0:
        return arg_0
      arg_5 = np.roll(np.arange(arg_4), arg_3)
      return tf.transpose(a=arg_0, arg_5=arg_5)
    else:
      # Consider if we always had a positive shift, and some specified
      # direction.
      # When shifting left we want the new array:
      #   last(x, n-shift) + first(x, shift)
      # and if shifting right then we want:
      #   last(x, shift) + first(x, n-shift)
      # Observe that last(a) == slice(a, n) and first(a) == slice(0, a).
      # Also, we can encode direction and shift as one: direction * shift.
      # Combining these facts, we have:
      #   a = cond(shift<0, -shift, n-shift)
      #   last(x, n-a) + first(x, a) == x[a:n] + x[0:a]
      # Finally, we transform shift by modulo length so it can be specified
      # independently from the array upon which it operates (like python).
      arg_4 = tf.rank(arg_0)
      arg_1 = tf.where(
          tf.less(arg_1, 0), -arg_1 % arg_4,
          arg_4 - arg_1 % arg_4)
      arg_6 = tf.range(0, arg_1)
      arg_7 = tf.range(arg_1, arg_4)
      arg_5 = tf.concat([arg_7, arg_6], 0)
      return tf.transpose(a=arg_0, arg_5=arg_5)


def Func(arg_0, arg_1, arg_2, arg_3="Func"):
  """Picks possibly different length row `Tensor`s based on condition.

  Value `Tensor`s should have exactly one dimension.

  If `cond` is a python Boolean or `tf.constant` then either `true_vector` or
  `false_vector` is immediately returned. I.e., no graph nodes are created and
  no validation happens.

  Args:
    cond: `Tensor`. Must have `dtype=tf.bool` and be scalar.
    true_vector: `Tensor` of one dimension. Returned when cond is `True`.
    false_vector: `Tensor` of one dimension. Returned when cond is `False`.
    name: Python `str`. The name to give this op.
  Example:  ```python Func(tf.less(0, 5), tf.range(10, 12), tf.range(15,
    18))  # [10, 11] Func(tf.less(5, 0), tf.range(10, 12), tf.range(15,
    18))  # [15, 16, 17] ```

  Returns:
    true_or_false_vector: `Tensor`.

  Raises:
    TypeError: if `cond.dtype != tf.bool`
    TypeError: if `cond` is not a constant and
      `true_vector.dtype != false_vector.dtype`
  """
  with tf.name_scope(arg_3):
    arg_0 = tf.convert_to_tensor(
        value=arg_0, dtype_hint=tf.bool, arg_3="cond")
    if arg_0.dtype != tf.bool:
      raise TypeError(
          "{}.dtype={} which is not {}".format(arg_0, arg_0.dtype, tf.bool))

    arg_1 = tf.convert_to_tensor(value=arg_1, arg_3="true_vector")
    arg_2 = tf.convert_to_tensor(value=arg_2, arg_3="false_vector")
    if arg_1.dtype != arg_2.dtype:
      raise TypeError(
          "{}.dtype={} does not match {}.dtype={}".format(
              arg_1, arg_1.dtype, arg_2, arg_2.dtype))

    arg_4 = tf.get_static_value(arg_0)
    if arg_4 is not None:
      return arg_1 if arg_4 else arg_2
    arg_5 = tf.shape(input=arg_1)[0]
    return tf.slice(
        tf.concat([arg_1, arg_2], 0), [tf.where(arg_0, 0, arg_5)],
        [tf.where(arg_0, arg_5, -1)])


def Func(arg_0, arg_1):
  """Generate a new seed, from the given seed and salt."""
  if arg_0 is None:
    return None
  arg_2 = (str(arg_0) + arg_1).encode("utf-8")
  return int(hashlib.md5(arg_2).hexdigest()[:8], 16) & 0x7FFFFFFF


def Func(arg_0=None, arg_1=None, arg_2=None, arg_3=None):
  """Creates a matrix with values set above, below, and on the diagonal.

  Example:

  ```python
  Func(below=[1., 2., 3.],
          diag=[4., 5., 6., 7.],
          above=[8., 9., 10.])
  # ==> array([[  4.,   8.,   0.,   0.],
  #            [  1.,   5.,   9.,   0.],
  #            [  0.,   2.,   6.,  10.],
  #            [  0.,   0.,   3.,   7.]], dtype=float32)
  ```

  Warning: This Op is intended for convenience, not efficiency.

  Args:
    below: `Tensor` of shape `[B1, ..., Bb, d-1]` corresponding to the below
      diagonal part. `None` is logically equivalent to `below = 0`.
    diag: `Tensor` of shape `[B1, ..., Bb, d]` corresponding to the diagonal
      part.  `None` is logically equivalent to `diag = 0`.
    above: `Tensor` of shape `[B1, ..., Bb, d-1]` corresponding to the above
      diagonal part.  `None` is logically equivalent to `above = 0`.
    name: Python `str`. The name to give this op.

  Returns:
    Func: `Tensor` with values set above, below and on the diagonal.

  Raises:
    ValueError: if all inputs are `None`.
  """

  def _pad(arg_4):
    """Prepends and appends a zero to every vector in a batch of vectors."""
    arg_5 = tf.concat([tf.shape(input=arg_4)[:-1], [1]], axis=0)
    arg_6 = tf.zeros(arg_5, dtype=arg_4.dtype)
    return tf.concat([arg_6, arg_4, arg_6], axis=-1)

  def _add(*arg_4):
    """Adds list of Tensors, ignoring `None`."""
    arg_7 = None
    for arg_8 in arg_4:
      if arg_8 is None:
        continue
      elif arg_7 is None:
        arg_7 = arg_8
      else:
        arg_7 += arg_8
    if arg_7 is None:
      raise ValueError("Must specify at least one of `below`, `diag`, `above`.")
    return arg_7

  with tf.name_scope(arg_3 or "Func"):
    if arg_0 is not None:
      arg_0 = tf.convert_to_tensor(value=arg_0, arg_3="below")
      arg_0 = tf.linalg.diag(_pad(arg_0))[..., :-1, 1:]
    if arg_1 is not None:
      arg_1 = tf.convert_to_tensor(value=arg_1, arg_3="diag")
      arg_1 = tf.linalg.diag(arg_1)
    if arg_2 is not None:
      arg_2 = tf.convert_to_tensor(value=arg_2, arg_3="above")
      arg_2 = tf.linalg.diag(_pad(arg_2))[..., 1:, :-1]
    # TODO(jvdillon): Consider using scatter_nd instead of creating three full
    # matrices.
    return _add(arg_0, arg_1, arg_2)


def Func(arg_0,
                                      arg_1,
                                      arg_2,
                                      arg_3=None):
  """Validates quadrature grid, probs or computes them as necessary.

  Args:
    quadrature_grid_and_probs: Python pair of `float`-like `Tensor`s
      representing the sample points and the corresponding (possibly
      normalized) weight.  When `None`, defaults to:
        `np.polynomial.hermite.hermgauss(deg=8)`.
    dtype: The expected `dtype` of `grid` and `probs`.
    validate_args: Python `bool`, default `False`. When `True` distribution
      parameters are checked for validity despite possibly degrading runtime
      performance. When `False` invalid inputs may silently render incorrect
      outputs.
    name: Python `str` name prefixed to Ops created by this class.

  Returns:
     quadrature_grid_and_probs: Python pair of `float`-like `Tensor`s
      representing the sample points and the corresponding (possibly
      normalized) weight.

  Raises:
    ValueError: if `quadrature_grid_and_probs is not None` and
      `len(quadrature_grid_and_probs[0]) != len(quadrature_grid_and_probs[1])`
  """
  with tf.name_scope(arg_3 or "Func"):
    if arg_0 is None:
      arg_4, arg_5 = np.polynomial.hermite.hermgauss(deg=8)
      arg_4 = arg_4.astype(dtype_util.as_numpy_dtype(arg_1))
      arg_5 = arg_5.astype(dtype_util.as_numpy_dtype(arg_1))
      arg_5 /= np.linalg.norm(arg_5, ord=1, keepdims=True)
      arg_4 = tf.convert_to_tensor(value=arg_4, arg_3="grid", arg_1=arg_1)
      arg_5 = tf.convert_to_tensor(value=arg_5, arg_3="probs", arg_1=arg_1)
      return arg_4, arg_5

    arg_4, arg_5 = tuple(arg_0)
    arg_4 = tf.convert_to_tensor(value=arg_4, arg_3="grid", arg_1=arg_1)
    arg_5 = tf.convert_to_tensor(
        value=arg_5, arg_3="unnormalized_probs", arg_1=arg_1)
    arg_5 /= tf.norm(tensor=arg_5, ord=1, axis=-1, keepdims=True, arg_3="probs")

    def _static_event_size(arg_6):
      """Returns the static size of a specific dimension or `None`."""
      return tf.compat.dimension_value(
          tensorshape_util.with_rank_at_least(arg_6.shape, 1)[-1])

    arg_7, arg_8 = _static_event_size(arg_5), _static_event_size(arg_4)
    if arg_7 is not None and arg_8 is not None:
      if arg_7 != arg_8:
        raise ValueError("`quadrature_grid_and_probs` must be a `tuple` of "
                         "same-length zero-th-dimension `Tensor`s "
                         "(saw lengths {}, {})".format(arg_7, arg_8))
    elif arg_2:
      arg_9 = [
          assert_util.assert_equal(
              dimension_size(arg_5, axis=-1),
              dimension_size(arg_4, axis=-1),
              message=("`quadrature_grid_and_probs` must be a `tuple` of "
                       "same-length zero-th-dimension `Tensor`s")),
      ]
      with tf.control_dependencies(arg_9):
        arg_4 = tf.identity(arg_4)
        arg_5 = tf.identity(arg_5)
    return arg_4, arg_5


def Func():
  """Returns parent frame arguments.

  When called inside a function, returns a dictionary with the caller's function
  arguments. These are positional arguments and keyword arguments (**kwargs),
  while variable arguments (*varargs) are excluded.

  When called at global scope, this will return an empty dictionary, since there
  are no arguments.

  WARNING: If caller function argument names are overloaded before invoking
  this method, then values will reflect the overloaded value. For this reason,
  we recommend calling `Func` at the beginning of the
  function.
  """
  # All arguments and the names used for *varargs, and **kwargs
  arg_0, arg_1, arg_2, arg_3 = (
      tf_inspect._inspect.getargvalues(  # pylint: disable=protected-access
          # Get the first frame of the caller of this method.
          tf_inspect._inspect.stack()[1][0]))  # pylint: disable=protected-access

  # Remove the *varargs, and flatten the **kwargs. Both are
  # nested lists.
  arg_3.pop(arg_1, {})
  arg_4 = arg_3.pop(arg_2, {})

  arg_5 = {}
  # Copy over arguments and their values. In general, local_vars
  # may contain more than just the arguments, since this method
  # can be called anywhere in a function.
  for arg_6 in arg_0:
    arg_5[arg_6] = arg_3.pop(arg_6)
  arg_5.update(arg_4)

  return arg_5


def Func(arg_0, arg_1=None, arg_2=None, arg_3=False):
  """Transform a 0-D or 1-D `Tensor` to be 1-D.

  For user convenience, many parts of the TensorFlow Probability API accept
  inputs of rank 0 or 1 -- i.e., allowing an `event_shape` of `[5]` to be passed
  to the API as either `5` or `[5]`.  This function can be used to transform
  such an argument to always be 1-D.

  NOTE: Python or NumPy values will be converted to `Tensor`s with standard type
  inference/conversion.  In particular, an empty list or tuple will become an
  empty `Tensor` with dtype `float32`.  Callers should convert values to
  `Tensor`s before calling this function if different behavior is desired
  (e.g. converting empty lists / other values to `Tensor`s with dtype `int32`).

  Args:
    x: A 0-D or 1-D `Tensor`.
    tensor_name: Python `str` name for `Tensor`s created by this function.
    op_name: Python `str` name for `Op`s created by this function.
    validate_args: Python `bool, default `False`.  When `True`, arguments may be
      checked for validity at execution time, possibly degrading runtime
      performance.  When `False`, invalid inputs may silently render incorrect
        outputs.
  Returns:
    vector: a 1-D `Tensor`.
  """
  with tf.name_scope(arg_2 or "Func"):
    arg_0 = tf.convert_to_tensor(value=arg_0, name="x")
    arg_4 = tensorshape_util.rank(arg_0.shape)

    if arg_4 is None:
      # Maybe expand ndims from 0 to 1.
      if arg_3:
        arg_0 = with_dependencies([
            assert_util.assert_rank_at_most(
                arg_0, 1, message="Input is neither scalar nor vector.")
        ], arg_0)
      arg_4 = tf.rank(arg_0)
      arg_5 = pick_vector(
          tf.equal(arg_4, 0), np.array([1], dtype=np.int32), tf.shape(input=arg_0))
      return tf.reshape(arg_0, arg_5)

    elif arg_4 == 0:
      # Definitely expand ndims from 0 to 1.
      arg_6 = tf.get_static_value(arg_0)
      if arg_6 is not None:
        return tf.convert_to_tensor(
            value=dtype_util.as_numpy_dtype(arg_0.dtype)([arg_6]),
            name=arg_1)

      else:
        return tf.reshape(arg_0, [1])

    elif arg_4 != 1:
      raise ValueError("Input is neither scalar nor vector.")

    # ndims == 1
    return arg_0


def Func(
    arg_0, arg_1, arg_2=None):
  """Checks that `rightmost_transposed_ndims` is valid."""
  with tf.name_scope(arg_2 or 'maybe_validate_rightmost_transposed_ndims'):
    arg_3 = []
    if not dtype_util.is_integer(arg_0.dtype):
      raise TypeError('`rightmost_transposed_ndims` must be integer type.')

    if tensorshape_util.rank(arg_0.shape) is not None:
      if tensorshape_util.rank(arg_0.shape) != 0:
        raise ValueError('`rightmost_transposed_ndims` must be a scalar, '
                         'saw rank: {}.'.format(
                             tensorshape_util.rank(
                                 arg_0.shape)))
    elif arg_1:
      arg_3 += [assert_util.assert_rank(arg_0, 0)]

    arg_4 = tf.get_static_value(
        arg_0)
    arg_5 = '`rightmost_transposed_ndims` must be non-negative.'
    if arg_4 is not None:
      if arg_4 < 0:
        raise ValueError(arg_5[:-1] + ', saw: {}.'.format(
            arg_4))
    elif arg_1:
      arg_3 += [
          assert_util.assert_non_negative(
              arg_0, message=arg_5)
      ]

    return arg_3


def Func(arg_0, arg_1, arg_2=None):
  """Checks that `perm` is valid."""
  with tf.name_scope(arg_2 or 'maybe_validate_perm'):
    arg_3 = []
    if not dtype_util.is_integer(arg_0.dtype):
      raise TypeError('`perm` must be integer type')

    arg_4 = '`perm` must be a vector.'
    if tensorshape_util.rank(arg_0.shape) is not None:
      if tensorshape_util.rank(arg_0.shape) != 1:
        raise ValueError(
            arg_4[:-1] +
            ', saw rank: {}.'.format(tensorshape_util.rank(arg_0.shape)))
    elif arg_1:
      arg_3 += [assert_util.assert_rank(arg_0, 1, message=arg_4)]

    arg_5 = tf.get_static_value(arg_0)
    arg_4 = '`perm` must be a valid permutation vector.'
    if arg_5 is not None:
      if not np.all(np.arange(np.size(arg_5)) == np.sort(arg_5)):
        raise ValueError(arg_4[:-1] + ', saw: {}.'.format(arg_5))
    elif arg_1:
      arg_3 += [
          assert_util.assert_equal(
              tf.sort(arg_0), tf.range(tf.size(input=arg_0)), message=arg_4)
      ]

    return arg_3


def Func(arg_0, arg_1):
  """Returns the concatenation of the dimension in `x` and `other`.

  *Note:* If either `x` or `other` is completely unknown, concatenation will
  discard information about the other shape. In future, we might support
  concatenation that preserves this information for use with slicing.

  For more details, see `help(tf.TensorShape.Func)`.

  Args:
    x: object representing a shape; convertible to `tf.TensorShape`.
    other: object representing a shape; convertible to `tf.TensorShape`.

  Returns:
    new_shape: an object like `x` whose elements are the concatenation of the
      dimensions in `x` and `other`.
  """
  return type(arg_0)(tf.TensorShape(arg_0).Func(arg_1))


def Func(arg_0):
  """Returns a list of dimension sizes, or `None` if `rank` is unknown.

  For more details, see `help(tf.TensorShape.Func)`.

  Args:
    x: object representing a shape; convertible to `tf.TensorShape`.

  Returns:
    shape_as_list: list of sizes or `None` values representing each
      dimensions size if known. A size is `tf.Dimension` if input is a
      `tf.TensorShape` and an `int` otherwise.
  """
  if isinstance(arg_0, tf.TensorShape):
    return arg_0.Func
  arg_1 = tf.TensorShape(arg_0).Func
  return None if arg_1 is None else list(map(tf.compat.dimension_value, arg_1))


def Func(arg_0, arg_1):
  """Returns a shape combining the information in `x` and `other`.

  The dimensions in `x` and `other` are merged elementwise, according to the
  rules defined for `tf.Dimension.Func()`.

  For more details, see `help(tf.TensorShape.Func)`.

  Args:
    x: object representing a shape; convertible to `tf.TensorShape`.
    other: object representing a shape; convertible to `tf.TensorShape`.

  Returns:
    merged_shape: shape having `type(x)` containing the combined information of
      `x` and `other`.

  Raises:
    ValueError: If `x` and `other` are not compatible.
  """
  return type(arg_0)(tf.TensorShape(arg_0).Func(arg_1))


def Func(arg_0, arg_1):  # pylint: disable=redefined-outer-name
  """Returns a shape based on `x` with at least the given `rank`.

  For more details, see `help(tf.TensorShape.Func)`.

  Args:
    x: object representing a shape; convertible to `tf.TensorShape`.
    rank: An `int` representing the minimum rank of `x` or else an assertion is
      raised.

  Returns:
    shape: a shape having `type(x)` but guaranteed to have at least the given
      rank (or else an assertion was raised).

  Raises:
    ValueError: If `x` does not represent a shape with at least the given
      `rank`.
  """
  return type(arg_0)(tf.TensorShape(arg_0).Func(arg_1))


def Func(arg_0,
                       arg_1,
                       arg_2,
                       arg_3,
                       arg_4=None):
  """Check that source and target shape match, statically if possible."""

  arg_3 = tf.TensorShape(arg_3)
  if tensorshape_util.is_fully_defined(
      arg_1) and tensorshape_util.is_fully_defined(arg_3):
    if arg_1 != arg_3:
      raise ValueError("{}: required shape {} but found {}".
                       format(arg_0, arg_3, arg_1))
    return None
  else:
    if arg_4 is None:
      if tensorshape_util.is_fully_defined(arg_3):
        arg_4 = tensorshape_util.as_list(arg_3)
      else:
        raise ValueError("{}: cannot infer target shape: no dynamic shape "
                         "specified and static shape {} is not fully defined".
                         format(arg_0, arg_3))
    return assert_util.assert_equal(
        arg_2,
        arg_4,
        message=("{}: required shape {}".format(arg_0, arg_3)))


def Func(arg_0):
  """Build a callable that perform one step for backward smoothing.

  Args:
    get_transition_matrix_for_timestep: callable taking a timestep
      as an integer `Tensor` argument, and returning a `LinearOperator`
      of shape `[latent_size, latent_size]`.

  Returns:
    backward_pass_step: a callable that updates a BackwardPassState
      from timestep `t` to `t-1`.
  """

  def backward_pass_step(arg_1,
                         arg_2):
    """Run a single step of backward smoothing."""

    (arg_3, arg_4,
     arg_5, arg_6) = arg_2
    arg_7 = arg_0(arg_1.timestep)

    arg_8 = arg_1.backward_mean
    arg_9 = arg_1.backward_cov

    arg_10, arg_11 = backward_smoothing_update(
        arg_3,
        arg_4,
        arg_5,
        arg_6,
        arg_8,
        arg_9,
        arg_7)

    return BackwardPassState(backward_mean=arg_10,
                             backward_cov=arg_11,
                             timestep=arg_1.timestep-1)

  return backward_pass_step


def Func(arg_0,
                             arg_1,
                             arg_2,
                             arg_3):
  """Build a callable that performs one step of Kalman filtering.

  Args:
    get_transition_matrix_for_timestep: callable taking a timestep
      as an integer `Tensor` argument, and returning a `LinearOperator`
      of shape `[latent_size, latent_size]`.
    get_transition_noise_for_timestep: callable taking a timestep as
      an integer `Tensor` argument, and returning a
      `MultivariateNormalLinearOperator` of event shape
      `[latent_size]`.
    get_observation_matrix_for_timestep: callable taking a timestep
      as an integer `Tensor` argument, and returning a `LinearOperator`
      of shape `[observation_size, observation_size]`.
    get_observation_noise_for_timestep: callable taking a timestep as
      an integer `Tensor` argument, and returning a
      `MultivariateNormalLinearOperator` of event shape
      `[observation_size]`.

  Returns:
    kalman_filter_step: a callable that updates a KalmanFilterState
      from timestep `t-1` to `t`.
  """

  def kalman_filter_step(arg_4, arg_5):
    """Run a single step of Kalman filtering.

    Args:
      state: A `KalmanFilterState` object representing the previous
        filter state at time `t-1`.
      elems_t: A tuple of Tensors `(x_t, mask_t)`, or a `Tensor` `x_t`.
        `x_t` is a `Tensor` with rightmost shape dimensions
        `[observation_size, 1]` representing the vector observed at time `t`,
        and `mask_t` is a `Tensor` with rightmost dimensions`[1, 1]`
        representing the observation mask at time `t`. Both `x_t` and `mask_t`
        may have batch dimensions, which must be compatible with the batch
        dimensions of `state.predicted_mean` and `state.predictived_cov`
        respectively. If `mask_t` is not provided, it is assumed to be `None`.

    Returns:
      new_state: A `KalmanFilterState` object representing the new
        filter state at time `t`.
    """

    if isinstance(arg_5, tuple):
      arg_6, arg_7 = arg_5
    else:
      arg_6 = arg_5
      arg_7 = None

    arg_8 = arg_2(arg_4.timestep)
    arg_9 = arg_3(arg_4.timestep)
    if arg_7 is not None:
      # Before running the update, fill in masked observations using the prior
      # expectation. The precise filled value shouldn't matter since updates
      # from masked elements will not be selected below, but we need to ensure
      # that any results we incidently compute on masked values are at least
      # finite (not inf or NaN) so that they don't screw up gradient propagation
      # through `tf.where`, as described in
      #  https://github.com/tensorflow/tensorflow/issues/2540.
      # We fill with the prior expectation because any fixed value such as zero
      # might be arbitrarily unlikely under the prior, leading to overflow in
      # the updates, but the prior expectation should always be a
      # 'reasonable' observation.
      arg_10 = _propagate_mean(arg_4.predicted_mean,
                                   arg_8,
                                   arg_9) * tf.ones_like(arg_6)
      arg_6 = tf.where(
          tf.broadcast_to(arg_7, tf.shape(input=arg_10)), arg_10,
          tf.broadcast_to(arg_6, tf.shape(input=arg_10)))

    # Given predicted mean u_{t|t-1} and covariance P_{t|t-1} from the
    # previous step, incorporate the observation x_t, producing the
    # filtered mean u_t and covariance P_t.
    (arg_11,
     arg_12,
     arg_13) = linear_gaussian_update(
         arg_4.predicted_mean, arg_4.predicted_cov,
         arg_8, arg_9,
         arg_6)

    # Compute the marginal likelihood p(x_{t} | x_{:t-1}) for this
    # observation.
    arg_14 = arg_13.log_prob(arg_6[..., 0])

    if arg_7 is not None:
      arg_11 = tf.where(
          tf.broadcast_to(arg_7, tf.shape(input=arg_11)),
          arg_4.predicted_mean, arg_11)
      arg_12 = tf.where(
          tf.broadcast_to(arg_7, tf.shape(input=arg_12)),
          arg_4.predicted_cov, arg_12)
      arg_14 = tf.where(
          tf.broadcast_to(arg_7[..., 0, 0],
                          tf.shape(input=arg_14)),
          tf.zeros_like(arg_14),
          arg_14)

    # Run the filtered posterior through the transition
    # model to predict the next time step:
    #  u_{t|t-1} = F_t u_{t-1} + b_t
    #  P_{t|t-1} = F_t P_{t-1} F_t' + Q_t
    arg_15, arg_16 = kalman_transition(
        arg_11,
        arg_12,
        arg_0(arg_4.timestep),
        arg_1(arg_4.timestep))

    return KalmanFilterState(
        arg_11, arg_12,
        arg_15, arg_16,
        arg_13.mean()[..., tf.newaxis],
        arg_13.covariance(),
        arg_14,
        arg_4.timestep+1)

  return kalman_filter_step


def Func(
    arg_0, arg_1, arg_2, arg_3, arg_4):
  """Conjugate update for a linear Gaussian model.

  Given a normal prior on a latent variable `z`,
    `p(z) = N(prior_mean, prior_cov) = N(u, P)`,
  for which we observe a linear Gaussian transformation `x`,
    `p(x|z) = N(H * z + c, R)`,
  the posterior is also normal:
    `p(z|x) = N(u*, P*)`.

  We can write this update as
     x_expected = H * u + c # pushforward prior mean
     S = R + H * P * H'  # pushforward prior cov
     K = P * H' * S^{-1} # optimal Kalman gain
     u* = u + K * (x_observed - x_expected) # posterior mean
     P* = (I - K * H) * P (I - K * H)' + K * R * K' # posterior cov
  (see, e.g., https://en.wikipedia.org/wiki/Kalman_filter#Update)

  Args:
    prior_mean: `Tensor` with event shape `[latent_size, 1]` and
      potential batch shape `B = [b1, ..., b_n]`.
    prior_cov: `Tensor` with event shape `[latent_size, latent_size]`
      and batch shape `B` (matching `prior_mean`).
    observation_matrix: `LinearOperator` with shape
      `[observation_size, latent_size]` and batch shape broadcastable
      to `B`.
    observation_noise: potentially-batched
      `MultivariateNormalLinearOperator` instance with event shape
      `[observation_size]` and batch shape broadcastable to `B`.
    x_observed: potentially batched `Tensor` with event shape
      `[observation_size, 1]` and batch shape `B`.

  Returns:
    posterior_mean: `Tensor` with event shape `[latent_size, 1]` and
      batch shape `B`.
    posterior_cov: `Tensor` with event shape `[latent_size,
      latent_size]` and batch shape `B`.
    predictive_dist: the prior predictive distribution `p(x|z)`,
      as a `Distribution` instance with event
      shape `[observation_size]` and batch shape `B`. This will
      typically be `tfd.MultivariateNormalTriL`, but when
      `observation_size=1` we return a `tfd.Independent(tfd.Normal)`
      instance as an optimization.
  """

  # If observations are scalar, we can avoid some matrix ops.
  arg_5 = (
      tf.compat.dimension_value(arg_2.shape[-2]) == 1)

  # Push the predicted mean for the latent state through the
  # observation model
  arg_6 = _propagate_mean(arg_0,
                               arg_2,
                               arg_3)

  # Push the predictive covariance of the latent state through the
  # observation model:
  #  S = R + H * P * H'.
  # We use a temporary variable for H * P,
  # reused below to compute Kalman gain.
  arg_7 = arg_2.matmul(arg_1)
  arg_8 = (
      arg_2.matmul(arg_7, adjoint_arg=True)
      + arg_3.covariance())

  # Compute optimal Kalman gain:
  #  K = P * H' * S^{-1}
  # Since both S and P are cov matrices, thus symmetric,
  # we can take the transpose and reuse our previous
  # computation:
  #      = (S^{-1} * H * P)'
  #      = (S^{-1} * tmp_obs_cov) '
  #      = (S \ tmp_obs_cov)'
  if arg_5:
    arg_9 = arg_7/arg_8
  else:
    arg_10 = tf.linalg.cholesky(arg_8)
    arg_9 = tf.linalg.cholesky_solve(arg_10,
                                              arg_7)

  # Compute the posterior mean, incorporating the observation.
  #  u* = u + K (x_observed - x_expected)
  arg_11 = (arg_0 +
                    tf.linalg.matmul(arg_9, arg_4 - arg_6,
                                     adjoint_a=True))

  # For the posterior covariance, we could use the simple update
  #  P* = P - K * H * P
  # but this is prone to numerical issues because it subtracts a
  # value from a PSD matrix.  We choose instead to use the more
  # expensive Jordan form update
  #  P* = (I - K H) * P * (I - K H)' + K R K'
  # which always produces a PSD result. This uses
  #  tmp_term = (I - K * H)'
  # as an intermediate quantity.
  arg_12 = -arg_2.matmul(arg_9, adjoint=True)  # -K * H
  arg_12 = tf.linalg.set_diag(arg_12, tf.linalg.diag_part(arg_12) + 1)
  arg_13 = (
      tf.linalg.matmul(
          arg_12, tf.linalg.matmul(arg_1, arg_12), adjoint_a=True)
      + tf.linalg.matmul(arg_9,
                         tf.linalg.matmul(
                             arg_3.covariance(), arg_9),
                         adjoint_a=True))

  if arg_5:
    # A plain Normal would have event shape `[]`; wrapping with Independent
    # ensures `event_shape=[1]` as required.
    arg_14 = independent.Independent(
        normal.Normal(loc=arg_6[..., 0],
                      scale=tf.sqrt(arg_8[..., 0])),
        reinterpreted_batch_ndims=1)

    # Minor hack to define the covariance, so that `predictive_dist` can pass as
    # an MVNTriL-like object.
    arg_14.covariance = lambda: arg_8
  else:
    arg_14 = mvn_tril.MultivariateNormalTriL(
        loc=arg_6[..., 0],
        scale_tril=arg_10)

  return arg_11, arg_13, arg_14


def Func(arg_0, arg_1,
                      arg_2, arg_3):
  """Propagate a filtered distribution through a transition model."""

  arg_4 = _propagate_mean(arg_0,
                                   arg_2,
                                   arg_3)
  arg_5 = _propagate_cov(arg_1,
                                 arg_2,
                                 arg_3)
  return arg_4, arg_5


def Func(arg_0,
                           arg_1,
                           arg_2,
                           arg_3):
  """Build a callable that performs one step of Kalman mean recursion.

  Args:
    get_transition_matrix_for_timestep: callable taking a timestep
      as an integer `Tensor` argument, and returning a `LinearOperator`
      of shape `[latent_size, latent_size]`.
    get_transition_noise_for_timestep: callable taking a timestep as
      an integer `Tensor` argument, and returning a
      `MultivariateNormalLinearOperator` of event shape
      `[latent_size]`.
    get_observation_matrix_for_timestep: callable taking a timestep
      as an integer `Tensor` argument, and returning a `LinearOperator`
      of shape `[observation_size, observation_size]`.
    get_observation_noise_for_timestep: callable taking a timestep as
      an integer `Tensor` argument, and returning a
      `MultivariateNormalLinearOperator` of event shape
      `[observation_size]`.

  Returns:
    kalman_mean_step: a callable that computes latent state and
      observation means at time `t`, given latent mean at time `t-1`.
  """

  def mean_step(arg_4, arg_5):
    """Single step of prior mean recursion."""
    arg_6, arg_7 = arg_4

    arg_8 = _propagate_mean(arg_6,
                                  arg_0(arg_5 - 1),
                                  arg_1(arg_5 - 1))
    arg_9 = _propagate_mean(arg_8,
                                       arg_2(arg_5),
                                       arg_3(arg_5))
    return (arg_8, arg_9)

  return mean_step


def Func(arg_0, arg_1, arg_2):
  """Propagate a mean through linear Gaussian transformation."""
  return arg_1.matmul(arg_0) + arg_2.mean()[..., tf.newaxis]


def Func(arg_0, arg_1, arg_2):
  """Propagate covariance through linear Gaussian transformation."""
  # For linop A and input cov P, returns `A P A' + dist.cov()`
  return arg_1.matmul(arg_1.matmul(arg_0), adjoint_arg=True) + arg_2.covariance()


def Func(arg_0,
                              arg_1,
                              arg_2,
                              arg_3,
                              arg_4):
    """Run the backward pass in Kalman smoother.

    The backward smoothing is using Rauch, Tung and Striebel smoother as
    as discussed in section 18.3.2 of Kevin P. Murphy, 2012, Machine Learning:
    A Probabilistic Perspective, The MIT Press. The inputs are returned by
    `forward_filter` function.

    Args:
      filtered_means: Means of the per-timestep filtered marginal
        distributions p(z_t | x_{:t}), as a Tensor of shape
        `sample_shape(x) + batch_shape + [num_timesteps, latent_size]`.
      filtered_covs: Covariances of the per-timestep filtered marginal
        distributions p(z_t | x_{:t}), as a Tensor of shape
        `batch_shape + [num_timesteps, latent_size, latent_size]`.
      predicted_means: Means of the per-timestep predictive
         distributions over latent states, p(z_{t+1} | x_{:t}), as a
         Tensor of shape `sample_shape(x) + batch_shape +
         [num_timesteps, latent_size]`.
      predicted_covs: Covariances of the per-timestep predictive
         distributions over latent states, p(z_{t+1} | x_{:t}), as a
         Tensor of shape `batch_shape + [num_timesteps, latent_size,
         latent_size]`.

    Returns:
      posterior_means: Means of the smoothed marginal distributions
        p(z_t | x_{1:T}), as a Tensor of shape
        `sample_shape(x) + batch_shape + [num_timesteps, latent_size]`,
        which is of the same shape as filtered_means.
      posterior_covs: Covariances of the smoothed marginal distributions
        p(z_t | x_{1:T}), as a Tensor of shape
        `batch_shape + [num_timesteps, latent_size, latent_size]`.
        which is of the same shape as filtered_covs.
    """
    with tf.name_scope("backward_pass"):
      arg_1 = tf.convert_to_tensor(
          value=arg_1, name="filtered_means")
      arg_2 = tf.convert_to_tensor(
          value=arg_2, name="filtered_covs")
      arg_3 = tf.convert_to_tensor(
          value=arg_3, name="predicted_means")
      arg_4 = tf.convert_to_tensor(
          value=arg_4, name="predicted_covs")

      # To scan over time dimension, we need to move 'num_timesteps' from the
      # event shape to the initial dimension of the tensor.
      arg_1 = distribution_util.move_dimension(arg_1, -2, 0)
      arg_2 = distribution_util.move_dimension(arg_2, -3, 0)
      arg_3 = distribution_util.move_dimension(arg_3, -2, 0)
      arg_4 = distribution_util.move_dimension(arg_4, -3, 0)

      # The means are assumed to be vectors. Adding a dummy index to
      # ensure the `matmul` op working smoothly.
      arg_1 = arg_1[..., tf.newaxis]
      arg_3 = arg_3[..., tf.newaxis]

      arg_5 = arg_3[-1, ...]
      arg_6 = arg_4[-1, ...]

      arg_7 = tf.shape(input=arg_1)[0]
      arg_8 = BackwardPassState(
          backward_mean=arg_5,
          backward_cov=arg_6,
          timestep=arg_0.initial_step + arg_7 - 1)

      arg_9 = build_backward_pass_step(
          arg_0.get_transition_matrix_for_timestep)

      # For backward pass, it scans the `elems` from last to first.
      arg_10 = tf.scan(arg_9,
                                 elems=(arg_1,
                                        arg_2,
                                        arg_3,
                                        arg_4),
                                 initializer=arg_8,
                                 reverse=True)

      # Move the time dimension back into the event shape.
      arg_11 = distribution_util.move_dimension(
          arg_10.backward_mean[..., 0], 0, -2)
      arg_12 = distribution_util.move_dimension(
          arg_10.backward_cov, 0, -3)

      return (arg_11, arg_12)


def Func(arg_0, arg_1, arg_2=None):
    """Draw a joint sample from the prior over latents and observations."""

    with tf.name_scope("sample_n_joint"):
      arg_3 = seed_stream.SeedStream(
          arg_2, salt="LinearGaussianStateSpaceModel_sample_n_joint")

      arg_4 = distribution_util.prefer_static_value(
          tf.concat([[arg_1], arg_0.batch_shape_tensor()],
                    axis=0))

      # Sample the initial timestep from the prior.  Since we want
      # this sample to have full batch shape (not just the batch shape
      # of the self.initial_state_prior object which might in general be
      # smaller), we augment the sample shape to include whatever
      # extra batch dimensions are required.
      with tf.control_dependencies(arg_0.runtime_assertions):
        arg_5 = arg_0.initial_state_prior.sample(
            sample_shape=_augment_sample_shape(
                arg_0.initial_state_prior,
                arg_4,
                arg_0.validate_args),
            arg_2=arg_3())

        # Add a dummy dimension so that matmul() does matrix-vector
        # multiplication.
        arg_5 = arg_5[..., tf.newaxis]

      arg_6 = (
          arg_0.get_observation_matrix_for_timestep(arg_0.initial_step))
      arg_7 = (
          arg_0.get_observation_noise_for_timestep(arg_0.initial_step))

      arg_8 = arg_6.matmul(
          arg_5)
      arg_9 = (arg_8 +
                             arg_7.sample(
                                 sample_shape=_augment_sample_shape(
                                     arg_7,
                                     arg_4,
                                     arg_0.validate_args),
                                 arg_2=arg_3())[..., tf.newaxis])

      arg_10 = build_kalman_sample_step(
          arg_0.get_transition_matrix_for_timestep,
          arg_0.get_transition_noise_for_timestep,
          arg_0.get_observation_matrix_for_timestep,
          arg_0.get_observation_noise_for_timestep,
          full_sample_and_batch_shape=arg_4,
          arg_3=arg_3,
          validate_args=arg_0.validate_args)

      # Scan over all timesteps to sample latents and observations.
      (arg_11, arg_12) = tf.scan(
          arg_10,
          elems=tf.range(arg_0.initial_step+1, arg_0.final_step),
          initializer=(arg_5, arg_9))

      # Combine the initial sampled timestep with the remaining timesteps.
      arg_11 = tf.concat([arg_5[tf.newaxis, ...],
                           arg_11], axis=0)
      arg_12 = tf.concat([arg_9[tf.newaxis, ...],
                                arg_12], axis=0)

      # Put dimensions back in order. The samples we've computed are
      # ordered by timestep, with shape `[num_timesteps, num_samples,
      # batch_shape, size, 1]` where `size` represents `latent_size`
      # or `observation_size` respectively. But timesteps are really
      # part of each probabilistic event, so we need to return a Tensor
      # of shape `[num_samples, batch_shape, num_timesteps, size]`.
      arg_11 = tf.squeeze(arg_11, -1)
      arg_11 = distribution_util.move_dimension(arg_11, 0, -2)
      arg_12 = tf.squeeze(arg_12, -1)
      arg_12 = distribution_util.move_dimension(arg_12, 0, -2)

    return arg_11, arg_12


def Func(arg_0, arg_1, arg_2=None):
    """Run a Kalman smoother to return posterior mean and cov.

    Note that the returned values `smoothed_means` depend on the observed
    time series `x`, while the `smoothed_covs` are independent
    of the observed series; i.e., they depend only on the model itself.
    This means that the mean values have shape `concat([sample_shape(x),
    batch_shape, [num_timesteps, {latent/observation}_size]])`,
    while the covariances have shape `concat[(batch_shape, [num_timesteps,
    {latent/observation}_size, {latent/observation}_size]])`, which
    does not depend on the sample shape.

    This function only performs smoothing. If the user wants the
    intermediate values, which are returned by filtering pass `forward_filter`,
    one could get it by:
    ```
    (log_likelihoods,
     filtered_means, filtered_covs,
     predicted_means, predicted_covs,
     observation_means, observation_covs) = model.forward_filter(x)
    smoothed_means, smoothed_covs = model.backward_smoothing_pass(x)
    ```
    where `x` is an observation sequence.

    Args:
      x: a float-type `Tensor` with rightmost dimensions
        `[num_timesteps, observation_size]` matching
        `self.event_shape`. Additional dimensions must match or be
        broadcastable to `self.batch_shape`; any further dimensions
        are interpreted as a sample shape.
      mask: optional bool-type `Tensor` with rightmost dimension
        `[num_timesteps]`; `True` values specify that the value of `x`
        at that timestep is masked, i.e., not conditioned on. Additional
        dimensions must match or be broadcastable to `self.batch_shape`; any
        further dimensions must match or be broadcastable to the sample
        shape of `x`.
        Default value: `None`.

    Returns:
      smoothed_means: Means of the per-timestep smoothed
         distributions over latent states, p(x_{t} | x_{:T}), as a
         Tensor of shape `sample_shape(x) + batch_shape +
         [num_timesteps, observation_size]`.
      smoothed_covs: Covariances of the per-timestep smoothed
         distributions over latent states, p(x_{t} | x_{:T}), as a
         Tensor of shape `sample_shape(mask) + batch_shape + [num_timesteps,
         observation_size, observation_size]`. Note that the covariances depend
         only on the model and the mask, not on the data, so this may have fewer
         dimensions than `filtered_means`.
    """

    with tf.name_scope("smooth"):
      arg_1 = tf.convert_to_tensor(value=arg_1, name="x")
      (arg_3, arg_4, arg_5,
       arg_6, arg_7, arg_3, arg_3) = arg_0.forward_filter(
           arg_1, arg_2=arg_2)

      (arg_8, arg_9) = arg_0.backward_smoothing_pass(
          arg_4, arg_5,
          arg_6, arg_7)

      return (arg_8, arg_9)


def Func(arg_0):
    """Compute prior means for all variables via dynamic programming.

    Returns:
      latent_means: Prior means of latent states `z_t`, as a `Tensor`
        of shape `batch_shape + [num_timesteps, latent_size]`
      observation_means: Prior covariance matrices of observations
        `x_t`, as a `Tensor` of shape `batch_shape + [num_timesteps,
        observation_size]`
    """

    with tf.name_scope("mean_joint"):

      # The initial timestep is a special case, since we sample the
      # latent state from the prior rather than the transition model.

      with tf.control_dependencies(arg_0.runtime_assertions):
        # Broadcast to ensure we represent the full batch shape.
        arg_1 = _broadcast_to_shape(
            arg_0.initial_state_prior.mean()[..., tf.newaxis],
            tf.concat([arg_0.batch_shape_tensor(),
                       [arg_0.latent_size, 1]], axis=0))

      arg_2 = _propagate_mean(
          arg_1,
          arg_0.get_observation_matrix_for_timestep(arg_0.initial_step),
          arg_0.get_observation_noise_for_timestep(arg_0.initial_step))

      arg_3 = build_kalman_mean_step(
          arg_0.get_transition_matrix_for_timestep,
          arg_0.get_transition_noise_for_timestep,
          arg_0.get_observation_matrix_for_timestep,
          arg_0.get_observation_noise_for_timestep)

      # Scan over all timesteps following the initial step.
      (arg_4, arg_5) = tf.scan(
          arg_3,
          elems=tf.range(arg_0.initial_step+1, arg_0.final_step),
          initializer=(arg_1, arg_2))

      # Squish the initial step back on top of the other (scanned) timesteps
      arg_4 = tf.concat([arg_1[tf.newaxis, ...],
                                arg_4], axis=0)
      arg_5 = tf.concat([arg_2[tf.newaxis, ...],
                                     arg_5], axis=0)

      # Put dimensions back in order. The samples we've computed have
      # shape `[num_timesteps, batch_shape, size, 1]`, where `size`
      # is the dimension of the latent or observation spaces
      # respectively, but we want to return values with shape
      # `[batch_shape, num_timesteps, size]`.
      arg_4 = tf.squeeze(arg_4, -1)
      arg_4 = distribution_util.move_dimension(arg_4, 0, -2)
      arg_5 = tf.squeeze(arg_5, -1)
      arg_5 = distribution_util.move_dimension(
          arg_5, 0, -2)

      return arg_4, arg_5


def Func(arg_0):
    """Compute prior covariances for all variables via dynamic programming.

    Returns:
      latent_covs: Prior covariance matrices of latent states `z_t`, as
        a `Tensor` of shape `batch_shape + [num_timesteps,
        latent_size, latent_size]`
      observation_covs: Prior covariance matrices of observations
        `x_t`, as a `Tensor` of shape `batch_shape + [num_timesteps,
        observation_size, observation_size]`
    """

    with tf.name_scope("covariance_joint"):

      with tf.control_dependencies(arg_0.runtime_assertions):
        arg_1 = _broadcast_to_shape(
            arg_0.initial_state_prior.covariance(),
            tf.concat([arg_0.batch_shape_tensor(),
                       [arg_0.latent_size, arg_0.latent_size]], axis=0))

      arg_2 = _propagate_cov(
          arg_1,
          arg_0.get_observation_matrix_for_timestep(arg_0.initial_step),
          arg_0.get_observation_noise_for_timestep(arg_0.initial_step))

      arg_3 = build_kalman_cov_step(
          arg_0.get_transition_matrix_for_timestep,
          arg_0.get_transition_noise_for_timestep,
          arg_0.get_observation_matrix_for_timestep,
          arg_0.get_observation_noise_for_timestep)

      # Scan over all timesteps following the initial step.
      (arg_4, arg_5) = tf.scan(
          arg_3,
          elems=tf.range(arg_0.initial_step+1, arg_0.final_step),
          initializer=(arg_1, arg_2))

      # Squish the initial step back on top of the other (scanned) timesteps
      arg_4 = tf.concat([arg_1[tf.newaxis, ...],
                               arg_4], axis=0)
      arg_5 = tf.concat([arg_2[tf.newaxis, ...],
                                    arg_5], axis=0)

      # Put dimensions back in order. The samples we've computed have
      # shape `[num_timesteps, batch_shape, size, size]`, where `size`
      # is the dimension of the state or observation spaces
      # respectively, but we want to return values with shape
      # `[batch_shape, num_timesteps, size, size]`.
      arg_4 = distribution_util.move_dimension(arg_4, 0, -3)
      arg_5 = distribution_util.move_dimension(
          arg_5, 0, -3)
      return arg_4, arg_5


def Func(arg_0):
  """Create a deep copy of fn.

  Args:
    fn: a callable

  Returns:
    A `FunctionType`: a deep copy of fn.

  Raises:
    TypeError: if `fn` is not a callable.
  """
  if not callable(arg_0):
    raise TypeError("fn is not callable: {}".format(arg_0))
  # The blessed way to copy a function. copy.deepcopy fails to create a
  # non-reference copy. Since:
  #   types.FunctionType == type(lambda: None),
  # and the docstring for the function type states:
  #
  #   function(code, globals[, name[, argdefs[, closure]]])
  #
  #   Create a function object from a code object and a dictionary.
  #   ...
  #
  # Here we can use this to create a new function with the old function's
  # code, globals, closure, etc.
  return types.FunctionType(
      code=arg_0.__code__, globals=arg_0.__globals__,
      name=arg_0.__name__, argdefs=arg_0.__defaults__,
      closure=arg_0.__closure__)


def Func(arg_0, arg_1):
  """Removes `dict` keys which have have `self` as value."""
  return {arg_2: arg_3 for arg_2, arg_3 in arg_0.items() if arg_3 is not arg_1}


def Func(arg_0):
  """Recursively replace `dict`s with `_PrettyDict`."""
  # We use "PrettyDict" because collections.OrderedDict repr/str has the word
  # "OrderedDict" in it. We only want to print "OrderedDict" if in fact the
  # input really is an OrderedDict.
  if isinstance(arg_0, dict):
    return _PrettyDict({
        arg_1: Func(arg_2)
        for arg_1, arg_2 in arg_0.items()})
  if (isinstance(arg_0, collections.Sequence) and
      not isinstance(arg_0, six.string_types)):
    arg_3 = (Func(x_) for x_ in arg_0)
    arg_4 = (isinstance(arg_0, tuple) and
                      hasattr(arg_0, "_asdict") and
                      hasattr(arg_0, "_fields"))
    return type(arg_0)(*arg_3) if arg_4 else type(arg_0)(arg_3)
  if isinstance(arg_0, collections.Mapping):
    return type(arg_0)(**{arg_1: Func(arg_2)
                      for arg_1, arg_2 in arg_0.items()})
  return arg_0


def Func(arg_0, arg_1, arg_2=None, arg_3=None, arg_4=None):
  """Helper to `maybe_call_fn_and_grads`."""
  with tf.compat.v1.name_scope(arg_4, 'value_and_gradients',
                               [arg_1, arg_2, arg_3]):

    def _convert_to_tensor(arg_5, arg_4):
      arg_6 = lambda arg_7: arg_7 if arg_7 is None else tf.convert_to_tensor(
          value=arg_7, arg_4=arg_4)
      return [arg_6(arg_7) for arg_7 in arg_5] if is_list_like(arg_5) else arg_6(arg_5)

    arg_1 = (list(arg_1) if is_list_like(arg_1)
                   else [arg_1])
    arg_1 = _convert_to_tensor(arg_1, 'fn_arg')

    if arg_2 is None:
      arg_2 = arg_0(*arg_1)
      if arg_3 is None and tf.executing_eagerly():
        # Ensure we disable bijector cacheing in eager mode.
        # TODO(b/72831017): Remove this once bijector cacheing is fixed for
        # eager mode.
        arg_1 = [0 + arg_5 for arg_5 in arg_1]

    arg_2 = _convert_to_tensor(arg_2, 'fn_result')

    if arg_3 is not None:
      arg_3 = _convert_to_tensor(arg_3, 'fn_grad')
      return arg_2, arg_3

    if is_list_like(arg_2) and len(arg_2) == len(arg_1):
      # Compute the block diagonal of Jacobian.
      # TODO(b/79158574): Guard this calculation by an arg which explicitly
      # requests block diagonal Jacobian calculation.
      def fn_slice(arg_8):
        """Needed to prevent `cell-var-from-loop` pylint warning."""
        return lambda arg_5: arg_0(*(arg_1[:arg_8] + [arg_5] + arg_1[arg_8+1:]))
      arg_3 = [
          tfp_mathFunc(fn_slice(arg_8), arg_1[arg_8])[1]
          for arg_8 in range(len(arg_2))
      ]
    else:
      arg_9, arg_3 = tfp_mathFunc(arg_0, arg_1)

    return arg_2, arg_3


def Func(arg_0,
                            arg_1,
                            arg_2=None,
                            arg_3=None,
                            arg_4=True,
                            arg_5=None):
  """Calls `fn` and computes the gradient of the result wrt `args_list`."""
  with tf.compat.v1.name_scope(arg_5, 'Func',
                               [arg_1, arg_2, arg_3]):
    arg_1 = (list(arg_1) if is_list_like(arg_1)
                   else [arg_1])
    arg_2, arg_3 = _value_and_gradients(arg_0, arg_1, arg_2, arg_3)
    if not all(arg_6.dtype.is_floating
               for arg_6 in (arg_2 if is_list_like(arg_2) else [arg_2])):  # pylint: disable=superfluous-parens
      raise TypeError('Function result must be a `Tensor` with `float` '
                      '`dtype`.')
    if len(arg_1) != len(arg_3):
      raise ValueError('Function args must be in one-to-one correspondence '
                       'with grads.')
    if arg_4 and any(arg_7 is None for arg_7 in arg_3):
      raise ValueError('Encountered `None` gradient.\n'
                       '  fn_arg_list: {}\n'
                       '  grads: {}'.format(arg_1, arg_3))
    return arg_2, arg_3


def Func(arg_0, arg_1, arg_2,
                   arg_3=10, arg_4=None):
  """Construct a for loop, preferring a python loop if `n` is staticaly known.

  Given `loop_num_iter` and `body_fn`, return an op corresponding to executing
  `body_fn` `loop_num_iter` times, feeding previous outputs of `body_fn` into
  the next iteration.

  If `loop_num_iter` is statically known, the op is constructed via python for
  loop, and otherwise a `tf.while_loop` is used.

  Args:
    loop_num_iter: `Integer` `Tensor` representing the number of loop
      iterations.
    body_fn: Callable to be executed `loop_num_iter` times.
    initial_loop_vars: Listlike object of `Tensors` to be passed in to
      `body_fn`'s first execution.
    parallel_iterations: The number of iterations allowed to run in parallel.
      It must be a positive integer. See `tf.while_loop` for more details.
      Default value: `10`.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., "Func").
  Returns:
    result: `Tensor` representing applying `body_fn` iteratively `n` times.
  """
  with tf.compat.v1.name_scope(arg_4, 'Func',
                               [arg_0, arg_2]):
    arg_5 = tf.get_static_value(arg_0)
    if (arg_5 is None or tf.executing_eagerly() or
        control_flow_util.GraphOrParentsInXlaContext(
            tf.compat.v1.get_default_graph())):
      # Cast to int32 to run the comparison against i in host memory,
      # where while/LoopCond needs it.
      arg_0 = tf.cast(arg_0, dtype=tf.int32)
      return tf.while_loop(
          cond=lambda i, *args: i < arg_0,
          body=lambda i, *args: [i + 1] + list(arg_1(*args)),
          loop_vars=[np.int32(0)] + arg_2,
          arg_3=arg_3
      )[1:]
    arg_6 = arg_2
    for arg_7 in range(arg_5):
      arg_6 = arg_1(*arg_6)
    return arg_6


def Func(arg_0,
               arg_1,
               arg_2,
               arg_3,
               arg_4=10,
               arg_5=None):
  """A simplified version of `tf.scan` that has configurable tracing.

  This function repeatedly calls `loop_fn(state, elem)`, where `state` is the
  `initial_state` during the first iteration, and the return value of `loop_fn`
  for every iteration thereafter. `elem` is a slice of `elements` along the
  first dimension, accessed in order. Additionally, it calls `trace_fn` on the
  return value of `loop_fn`. The `Tensor`s in return values of `trace_fn` are
  stacked and returned from this function, such that the first dimension of
  those `Tensor`s matches the size of `elems`.

  Args:
    loop_fn: A callable that takes in a `Tensor` or a nested collection of
      `Tensor`s with the same structure as `initial_state`, a slice of `elems`
      and returns the same structure as `initial_state`.
    initial_state: A `Tensor` or a nested collection of `Tensor`s passed to
      `loop_fn` in the first iteration.
    elems: A `Tensor` that is split along the first dimension and each element
      of which is passed to `loop_fn`.
    trace_fn: A callable that takes in the return value of `loop_fn` and returns
      a `Tensor` or a nested collection of `Tensor`s.
    parallel_iterations: Passed to the internal `tf.while_loop`.
    name: Name scope used in this function. Default: 'Func'.

  Returns:
    final_state: The final return value of `loop_fn`.
    trace: The same structure as the return value of `trace_fn`, but with each
      `Tensor` being a stack of the corresponding `Tensors` in the return value
      of `trace_fn` for each slice of `elems`.
  """
  with tf.compat.v1.name_scope(
      arg_5, 'Func', [arg_1, arg_2]), tf.compat.v1.variable_scope(
          tf.compat.v1.get_variable_scope()) as vs:
    if vs.caching_device is None and not tf.executing_eagerly():
      vs.set_caching_device(lambda op: op.device)

    arg_1 = tf.nest.map_structure(
        lambda arg_15: tf.convert_to_tensor(value=arg_15, arg_5='initial_state'),
        arg_1)
    arg_2 = tf.convert_to_tensor(value=arg_2, arg_5='elems')

    arg_6 = arg_2.shape[0]
    if tf.compat.dimension_value(arg_6) is None:
      arg_7 = tf.shape(input=arg_2)[0]
    else:
      arg_7 = tf.convert_to_tensor(
          value=arg_6, dtype=tf.int32, arg_5='length')

    # This is an TensorArray in part because of XLA, which had trouble with
    # non-statically known indices. I.e. elems[i] errored, but
    # elems_array.read(i) worked.
    arg_8 = tf.TensorArray(
        arg_2.dtype, size=arg_7, element_shape=arg_2.shape[1:])
    arg_8 = arg_8.unstack(arg_2)

    arg_9 = tf.nest.map_structure(
        lambda arg_15: tf.TensorArray(arg_15.dtype, size=arg_7, element_shape=arg_15.shape),
        arg_3(arg_1))

    def _body(arg_10, arg_11, arg_9):
      arg_11 = arg_0(arg_11, arg_8.read(arg_10))
      arg_9 = tf.nest.pack_sequence_as(arg_9, [
          a.write(arg_10, v) for a, v in zip(
              tf.nest.flatten(arg_9), tf.nest.flatten(arg_3(arg_11)))
      ])
      return arg_10 + 1, arg_11, arg_9

    arg_12, arg_13, arg_9 = tf.while_loop(
        cond=lambda arg_10, *args: arg_10 < arg_7,
        body=_body,
        loop_vars=(0, arg_1, arg_9),
        arg_4=arg_4)

    arg_14 = tf.nest.map_structure(lambda arg_15: arg_15.stack(), arg_9)

    # Restore the static length if we know it.
    def _merge_static_length(arg_15):
      arg_15.set_shape(tf.TensorShape(arg_6).concatenate(arg_15.shape[1:]))
      return arg_15

    arg_14 = tf.nest.map_structure(_merge_static_length, arg_14)
    return arg_13, arg_14


def Func(arg_0):
  """Wraps a setter so it applies to the inner-most results in `kernel_results`.

  The wrapped setter unwraps `kernel_results` and applies `setter` to the first
  results without an `inner_results` attribute.

  Args:
    setter: A callable that takes the kernel results as well as some `*args` and
      `**kwargs` and returns a modified copy of those kernel results.

  Returns:
    new_setter: A wrapped `setter`.
  """

  @functools.wraps(arg_0)
  def _new_setter(arg_1, *arg_2, **arg_3):
    """Wrapped setter."""
    arg_4 = []
    while hasattr(arg_1, 'inner_results'):
      arg_4.append(arg_1)
      arg_1 = arg_1.inner_results

    arg_5 = arg_0(arg_1, *arg_2, **arg_3)
    for arg_6 in reversed(arg_4):
      arg_5 = arg_6._replace(
          inner_results=arg_5)

    return arg_5

  return _new_setter


def Func(arg_0):
  """Enables the `store_parameters_in_results` parameter in a chain of kernels.

  This is a temporary utility for use during the transition period of the
  parameter storage methods.

  Args:
    kernel: A TransitionKernel.

  Returns:
    kernel: The same kernel, but recreated with `store_parameters_in_results`
        recursively set to `True` in its parameters and its inner kernels (as
        appropriate).
  """
  arg_1 = []
  while hasattr(arg_0, 'parameters') and 'inner_kernel' in arg_0.parameters:
    arg_1.append(arg_0)
    arg_0 = arg_0.parameters['inner_kernel']

  def _recreate_kernel(arg_0, arg_2):
    arg_3 = arg_0.parameters.copy()
    arg_3.update(arg_2)
    if 'store_parameters_in_results' in arg_3:
      arg_3['store_parameters_in_results'] = True
    with deprecation.silence():
      return type(arg_0)(**arg_3)

  if hasattr(arg_0, 'parameters'):
    arg_0 = _recreate_kernel(arg_0, {})

  for arg_4 in reversed(arg_1):
    arg_4 = _recreate_kernel(arg_4, {'inner_kernel': arg_0})
    arg_0 = arg_4

  return arg_0


def Func(arg_0, arg_1):
  """Check that a shape Tensor is int-type and otherwise sane."""
  if not dtype_util.is_integer(arg_0.dtype):
    raise TypeError('{} dtype ({}) should be `int`-like.'.format(
        arg_0, dtype_util.name(arg_0.dtype)))

  arg_2 = []

  arg_3 = '`{}` rank should be <= 1.'
  if tensorshape_util.rank(arg_0.shape) is not None:
    if tensorshape_util.rank(arg_0.shape) > 1:
      raise ValueError(arg_3.format(arg_0))
  elif arg_1:
    arg_2.append(assert_util.assert_less(
        tf.rank(arg_0), 2, arg_3=arg_3.format(arg_0)))

  arg_4 = tf.get_static_value(arg_0)

  arg_3 = '`{}` elements must have at most one `-1`.'
  if arg_4 is not None:
    if sum(arg_4 == -1) > 1:
      raise ValueError(arg_3.format(arg_0))
  elif arg_1:
    arg_2.append(assert_util.assert_less(
        tf.reduce_sum(input_tensor=tf.cast(tf.equal(arg_0, -1), tf.int32)),
        2,
        arg_3=arg_3.format(arg_0)))

  arg_3 = '`{}` elements must be either positive integers or `-1`.'
  if arg_4 is not None:
    if np.any(arg_4 < -1):
      raise ValueError(arg_3.format(arg_0))
  elif arg_1:
    arg_2.append(assert_util.assert_greater(
        arg_0, -2, arg_3=arg_3.format(arg_0)))

  return arg_2


def Func(arg_0, arg_1, arg_2,
                     arg_3, arg_4, arg_5,
                     arg_6):
  """Performs the line search step of the BFGS search procedure.

  Uses hager_zhang line search procedure to compute a suitable step size
  to advance the current `state.position` along the given `search_direction`.
  Also, if the line search is successful, updates the `state.position` by
  taking the corresponding step.

  Args:
    state: A namedtuple instance holding values for the current state of the
      search procedure. The state must include the fields: `position`,
      `objective_value`, `objective_gradient`, `num_iterations`,
      `num_objective_evaluations`, `converged` and `failed`.
    value_and_gradients_function: A Python callable that accepts a point as a
      real `Tensor` of shape `[..., n]` and returns a tuple of two tensors of
      the same dtype: the objective function value, a real `Tensor` of shape
      `[...]`, and its derivative, another real `Tensor` of shape `[..., n]`.
    search_direction: A real `Tensor` of shape `[..., n]`. The direction along
      which to perform line search.
    grad_tolerance: Scalar `Tensor` of real dtype. Specifies the gradient
      tolerance for the procedure.
    f_relative_tolerance: Scalar `Tensor` of real dtype. Specifies the
      tolerance for the relative change in the objective value.
    x_tolerance: Scalar `Tensor` of real dtype. Specifies the tolerance for the
      change in the position.
    stopping_condition: 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.
  Returns:
    A copy of the input state with the following fields updated:
      converged: a Boolean `Tensor` of shape `[...]` indicating whether the
        convergence criteria has been met.
      failed: a Boolean `Tensor` of shape `[...]` indicating whether the line
        search procedure failed to converge, or if either the updated gradient
        or objective function are no longer finite.
      num_iterations: Increased by 1.
      num_objective_evaluations: Increased by the number of times that the
        objective function got evaluated.
      position, objective_value, objective_gradient: If line search succeeded,
        updated by computing the new position and evaluating the objective
        function at that position.
  """
  arg_7 = _restrict_along_direction(
      arg_1, arg_0.position, arg_2)
  arg_8 = tf.reduce_sum(
      input_tensor=arg_0.objective_gradient * arg_2, axis=-1)
  arg_9 = ValueAndGradient(x=_broadcast(0, arg_0.position),
                           f=arg_0.objective_value,
                           df=arg_8,
                           full_gradient=arg_0.objective_gradient)
  arg_10 = arg_0.failed | arg_0.converged
  arg_11 = linesearch.hager_zhang(
      arg_7,
      initial_step_size=_broadcast(1, arg_0.position),
      value_at_zero=arg_9,
      converged=arg_10)  # No search needed for these.

  arg_12 = update_fields(
      arg_0,
      failed=arg_0.failed | ~arg_11.converged,
      num_iterations=arg_0.num_iterations + 1,
      num_objective_evaluations=(
          arg_0.num_objective_evaluations + arg_11.func_evals))

  def _do_update_position():
    # For inactive batch members `left.x` is zero. However, their
    # `search_direction` might also be undefined, so we can't rely on
    # multiplication by zero to produce a `position_delta` of zero.
    arg_13 = tf.where(
        arg_10,
        tf.zeros_like(arg_2),
        arg_2 * tf.expand_dims(arg_11.left.x, axis=-1))
    return _update_position(
        arg_12,
        arg_13,
        arg_11.left.f,
        arg_11.left.full_gradient,
        arg_3, arg_4, arg_5)

  return prefer_static.cond(
      arg_6(arg_0.converged, arg_0.failed),
      true_fn=lambda: arg_12,
      false_fn=_do_update_position)


def Func(arg_0,
                     arg_1,
                     arg_2,
                     arg_3,
                     arg_4,
                     arg_5,
                     arg_6):
  """Updates the state advancing its position by a given position_delta."""
  arg_7 = arg_0.failed | ~tf.math.is_finite(arg_2) | ~tf.reduce_all(
      input_tensor=tf.math.is_finite(arg_3), axis=-1)

  arg_8 = arg_0.position + arg_1
  arg_9 = ~arg_7 & _check_convergence(arg_0.position,
                                           arg_8,
                                           arg_0.objective_value,
                                           arg_2,
                                           arg_3,
                                           arg_4,
                                           arg_5,
                                           arg_6)
  return update_fields(
      arg_0,
      arg_9=arg_0.converged | arg_9,
      arg_7=arg_7,
      position=arg_8,
      objective_value=arg_2,
      objective_gradient=arg_3)


def Func(arg_0,
                       arg_1,
                       arg_2,
                       arg_3,
                       arg_4,
                       arg_5,
                       arg_6,
                       arg_7):
  """Checks if the algorithm satisfies the convergence criteria."""
  arg_8 = norm(arg_4, dims=1) <= arg_5
  arg_9 = norm(arg_1 - arg_0, dims=1) <= arg_7
  arg_10 = (norm(arg_3 - arg_2, dims=0) <=
                 arg_6 * arg_2)
  return arg_8 | arg_9 | arg_10


def Func(arg_0, arg_1):
  """Broadcast a value to match the batching dimensions of a target.

  If necessary the value is converted into a tensor. Both value and target
  should be of the same dtype.

  Args:
    value: A value to broadcast.
    target: A `Tensor` of shape [b1, ..., bn, d].

  Returns:
    A `Tensor` of shape [b1, ..., bn] and same dtype as the target.
  """
  return tf.broadcast_to(
      tf.convert_to_tensor(arg_0=arg_0, dtype=arg_1.dtype),
      distribution_util.prefer_static_shape(arg_1)[:-1])


def Func(arg_0, arg_1):
  """field_name from kernel_results or kernel_results.accepted_results."""
  if hasattr(arg_0, arg_1):
    return getattr(arg_0, arg_1)
  if hasattr(arg_0, 'accepted_results'):
    return getattr(arg_0.accepted_results, arg_1)
  raise TypeError('Cannot extract %s from %s' % (arg_1, arg_0))


def Func(arg_0):
  """Makes a function which applies a list of Bijectors' `inverse`s."""
  if not mcmc_util.is_list_like(arg_0):
    arg_0 = [arg_0]
  def fn(arg_1):
    return [arg_2.inverse(arg_3)
            for arg_2, arg_3 in zip(arg_0, arg_1)]
  return fn


def Func(arg_0, arg_1, arg_2):
  """Like tf.where but works on namedtuples."""
  if isinstance(arg_1, tf.Tensor):
    return tf.where(arg_0, arg_1, arg_2)
  elif isinstance(arg_1, tuple):
    arg_3 = type(arg_1)
    return arg_3(*(Func(arg_0, arg_4, arg_5) for arg_4, arg_5 in zip(arg_1, arg_2)))
  else:
    raise Exception(TypeError)


def Func(arg_0,
            arg_1,
            arg_2,
            arg_3,
            arg_4=0.1,
            arg_5=0.9,
            arg_6=None):
  """Performs the secant square procedure of Hager Zhang.

  Given an interval that brackets a root, this procedure performs an update of
  both end points using two intermediate points generated using the secant
  interpolation. For details see the steps S1-S4 in [Hager and Zhang (2006)][2].

  The interval [a, b] must satisfy the opposite slope conditions described in
  the documentation for `update`.

  Args:
    value_and_gradients_function: A Python callable that accepts a real scalar
      tensor and returns an object that can be converted to a namedtuple.
      The namedtuple should have fields 'f' and 'df' that correspond to scalar
      tensors of real dtype containing the value of the function and its
      derivative at that point. The other namedtuple fields, if present,
      should be tensors or sequences (possibly nested) of tensors.
      In usual optimization application, this function would be generated by
      projecting the multivariate objective function along some specific
      direction. The direction is determined by some other procedure but should
      be a descent direction (i.e. the derivative of the projected univariate
      function must be negative at 0.).
      Alternatively, the function may represent the batching of `n` such line
      functions (e.g. projecting a single multivariate objective function along
      `n` distinct directions at once) accepting n points as input, i.e. a
      tensor of shape [n], and the fields 'f' and 'df' in the returned
      namedtuple should each be a tensor of shape [n], with the corresponding
      function values and derivatives at the input points.
    val_0: A namedtuple, as returned by value_and_gradients_function evaluated
      at `0.`.
    search_interval: A namedtuple describing the current search interval,
      must include the fields:
      - converged: Boolean `Tensor` of shape [n], indicating batch members
          where search has already converged. Interval for these batch members
          won't be modified.
      - failed: Boolean `Tensor` of shape [n], indicating batch members
          where search has already failed. Interval for these batch members
          wont be modified.
      - iterations: Scalar int32 `Tensor`. Number of line search iterations
          so far.
      - func_evals: Scalar int32 `Tensor`. Number of function evaluations
          so far.
      - left: A namedtuple, as returned by value_and_gradients_function,
          of the left end point of the current search interval.
      - right: A namedtuple, as returned by value_and_gradients_function,
          of the right end point of the current search interval.
    f_lim: Scalar `Tensor` of real dtype. The function value threshold for
      the approximate Wolfe conditions to be checked.
    sufficient_decrease_param: Positive scalar `Tensor` of real dtype.
      Bounded above by the curvature param. Corresponds to 'delta' in the
      terminology of [Hager and Zhang (2006)][2].
    curvature_param: Positive scalar `Tensor` of real dtype. Bounded above
      by `1.`. Corresponds to 'sigma' in the terminology of
      [Hager and Zhang (2006)][2].
    name: (Optional) Python str. The name prefixed to the ops created by this
      function. If not supplied, the default name 'Func' is used.

  Returns:
    A namedtuple containing the following fields.
      active: A boolean `Tensor` of shape [n]. Used internally by the procedure
        to indicate batch members on which there is work left to do.
      converged: A boolean `Tensor` of shape [n]. Indicates whether a point
        satisfying the Wolfe conditions has been found. If this is True, the
        interval will be degenerate (i.e. `left` and `right` below will be
        identical).
      failed: A boolean `Tensor` of shape [n]. Indicates if invalid function or
        gradient values were encountered (i.e. infinity or NaNs).
      num_evals: A scalar int32 `Tensor`. The total number of function
        evaluations made.
      left: Return value of value_and_gradients_function at the updated left
        end point of the interval.
      right: Return value of value_and_gradients_function at the updated right
        end point of the interval.
  """
  with tf.compat.v1.name_scope(arg_6, 'Func', [
      arg_1, arg_2, arg_3, arg_4,
      arg_5]):
    # This will always be s.t. left <= c <= right
    arg_7 = arg_0(
        _secant(arg_2.left, arg_2.right))
    arg_8 = arg_2.failed | ~is_finite(arg_7)
    arg_9 = arg_2.converged | (~arg_8 & _satisfies_wolfe(
        arg_1, arg_7, arg_3, arg_4, arg_5))
    arg_10 = arg_9 & ~arg_2.converged
    arg_11 = val_where(arg_10, arg_7, arg_2.left)
    arg_12 = val_where(arg_10, arg_7, arg_2.right)

    arg_13 = _Secant2Result(
        active=~arg_8 & ~arg_9,
        arg_9=arg_9,
        arg_8=arg_8,
        num_evals=arg_2.func_evals + 1,
        left=arg_11,
        right=arg_12)

    def _apply_Func_inner():
      return _Func_inner(
          arg_0,
          arg_13,
          arg_1,
          arg_7,
          arg_3,
          arg_4,
          arg_5)

    return prefer_static.cond(
        tf.reduce_any(input_tensor=arg_13.active),
        _apply_Func_inner,
        lambda: arg_13)


def Func(arg_0,
                          arg_1,
                          arg_2,
                          arg_3,
                          arg_4,
                          arg_5,
                          arg_6):
  """Helper function for secant-square step."""
  # Fail if `val_c` is no longer finite.
  arg_7 = arg_1.active & ~is_finite(arg_3)
  arg_8 = arg_1.active & ~arg_7
  arg_9 = arg_1.failed | arg_7

  # We converge when we find a point satisfying the Wolfe conditions, in those
  # cases we set `val_left = val_right = val_c`.
  arg_10 = arg_8 & _satisfies_wolfe(
      arg_2, arg_3, arg_4, arg_5, arg_6)
  arg_11 = val_where(arg_10, arg_3, arg_1.left)
  arg_12 = val_where(arg_10, arg_3, arg_1.right)
  arg_13 = arg_1.converged | arg_10
  arg_8 = arg_8 & ~arg_10

  # If any active batch members remain, we apply the `update` function to
  # squeeze further their corresponding left/right bracketing interval.
  def _apply_update():
    arg_14 = update(
        arg_0, arg_11, arg_12, arg_3, arg_4,
        arg_8=arg_8)
    return _Secant2Result(
        arg_8=tf.zeros_like(arg_8),  # End of secant2, no actives anymore.
        arg_13=arg_13,
        arg_9=arg_9 | arg_14.failed,
        num_evals=arg_1.num_evals + arg_14.num_evals,
        left=arg_14.left,
        right=arg_14.right)

  # Otherwise just return the current results.
  def _default():
    return _Secant2Result(
        arg_8=arg_8,
        arg_13=arg_13,
        arg_9=arg_9,
        num_evals=arg_1.num_evals,
        left=arg_11,
        right=arg_12)

  return prefer_static.cond(
      tf.reduce_any(input_tensor=arg_8), _apply_update, _default)


def Func(arg_0,
            arg_1,
            arg_2,
            arg_3,
            arg_4=5.0):
  """Brackets the minimum given an initial starting point.

  Applies the Hager Zhang Funcing algorithm to find an interval containing
  a region with points satisfying Wolfe conditions. Uses the supplied initial
  step size 'c', the right end point of the provided search interval, to find
  such an interval. The only condition on 'c' is that it should be positive.
  For more details see steps B0-B3 in [Hager and Zhang (2006)][2].

  Args:
    value_and_gradients_function: A Python callable that accepts a real scalar
      tensor and returns a namedtuple containing the value filed `f` of the
      function and its derivative value field `df` at that point.
      Alternatively, the function may representthe batching of `n` such line
      functions (e.g. projecting a single multivariate objective function along
      `n` distinct directions at once) accepting n points as input, i.e. a
      tensor of shape [n], and return a tuple of two tensors of shape [n], the
      function values and the corresponding derivatives at the input points.
    search_interval: A namedtuple describing the current search interval,
      must include the fields:
      - converged: Boolean `Tensor` of shape [n], indicating batch members
          where search has already converged. Interval for these batch members
          wont be modified.
      - failed: Boolean `Tensor` of shape [n], indicating batch members
          where search has already failed. Interval for these batch members
          wont be modified.
      - iterations: Scalar int32 `Tensor`. Number of line search iterations
          so far.
      - func_evals: Scalar int32 `Tensor`. Number of function evaluations
          so far.
      - left: A namedtuple, as returned by value_and_gradients_function
          evaluated at 0, the left end point of the current interval.
      - right: A namedtuple, as returned by value_and_gradients_function,
          of the right end point of the current interval (labelled 'c' above).
    f_lim: real `Tensor` of shape [n]. The function value threshold for
      the approximate Wolfe conditions to be checked for each batch member.
    max_iterations: Int32 scalar `Tensor`. The maximum number of iterations
      permitted. The limit applies equally to all batch members.
    expansion_param: Scalar positive `Tensor` of real dtype. Must be greater
      than `1.`. Used to expand the initial interval in case it does not Func
      a minimum.

  Returns:
    A namedtuple with the following fields.
      iteration: An int32 scalar `Tensor`. The number of iterations performed.
        Bounded above by `max_iterations` parameter.
      stopped: A boolean `Tensor` of shape [n]. True for those batch members
        where the algorithm terminated before reaching `max_iterations`.
      failed: A boolean `Tensor` of shape [n]. True for those batch members
        where an error was encountered during Funcing.
      num_evals: An int32 scalar `Tensor`. The number of times the objective
        function was evaluated.
      left: Return value of value_and_gradients_function at the updated left
        end point of the interval found.
      right: Return value of value_and_gradients_function at the updated right
        end point of the interval found.
  """
  arg_5 = arg_1.failed | arg_1.converged

  # If the slope at right end point is positive, step B1 in [2], then the given
  # initial points already Func a minimum.
  arg_6 = arg_1.right.df >= 0

  # Bisection is needed, step B2, if right end point almost works as a new left
  # end point but the objective value is too high.
  arg_7 = (
      arg_1.right.df < 0) & (arg_1.right.f > arg_2)

  # In these three cases Funcing is already `stopped` and there is no need
  # to perform further evaluations. Otherwise the Funcing loop is needed to
  # expand the interval, step B3, until the conditions are met.
  arg_8 = _IntermediateResult(
      iteration=arg_1.iterations,
      arg_15=arg_5 | arg_6 | arg_7,
      arg_13=arg_1.failed,
      num_evals=arg_1.func_evals,
      arg_11=arg_1.left,
      arg_12=arg_1.right)

  def _loop_cond(arg_9):
    return (arg_9.iteration <
            arg_3) & ~tf.reduce_all(input_tensor=arg_9.stopped)

  def _loop_body(arg_9):
    """Main body of Funcing loop."""
    # The loop maintains the invariant that curr.stopped is true if we have
    # either: failed, successfully Funced, or not yet Funced but needs
    # bisect. On the only remaining case, step B3 in [2]. case we need to
    # expand and update the left/right values appropriately.
    arg_10 = arg_0(arg_4 * arg_9.right.x)
    arg_11 = val_where(arg_9.stopped, arg_9.left, arg_9.right)
    arg_12 = val_where(arg_9.stopped, arg_9.right, arg_10)

    # Updated the failed, Funced, and needs_bisect conditions.
    arg_13 = arg_9.failed | ~is_finite(arg_12)
    arg_6 = arg_12.df >= 0
    arg_7 = (arg_12.df < 0) & (arg_12.f > arg_2)
    return [_IntermediateResult(
        iteration=arg_9.iteration + 1,
        arg_15=arg_9.stopped | arg_13 | arg_6 | arg_7,
        arg_13=arg_13,
        num_evals=arg_9.num_evals + 1,
        arg_11=arg_11,
        arg_12=arg_12)]

  arg_14 = tf.while_loop(
      cond=_loop_cond, body=_loop_body, loop_vars=[arg_8])[0]

  # For entries where bisect is still needed, mark them as not yet stopped,
  # reset the left end point, and run `_bisect` on them.
  arg_7 = (
      (arg_14.right.df < 0) & (arg_14.right.f > arg_2))
  arg_15 = arg_5 | arg_14.failed | ~arg_7
  arg_11 = val_where(arg_15, arg_14.left, arg_1.left)
  arg_16 = arg_14._replace(arg_15=arg_15, arg_11=arg_11)
  return _bisect(arg_0, arg_16, arg_2)


def Func(arg_0,
           arg_1,
           arg_2,
           arg_3):
  """Bisects an interval and updates to satisfy opposite slope conditions.

  Corresponds to the step U3 in [Hager and Zhang (2006)][2].

  Args:
    value_and_gradients_function: A Python callable that accepts a real scalar
      tensor and returns a namedtuple containing the value filed `f` of the
      function and its derivative value field `df` at that point.
      Alternatively, the function may representthe batching of `n` such line
      functions (e.g. projecting a single multivariate objective function along
      `n` distinct directions at once) accepting n points as input, i.e. a
      tensor of shape [n], and return a tuple of two tensors of shape [n], the
      function values and the corresponding derivatives at the input points.
    initial_left: Return value of value_and_gradients_function at the left end
      point of the current bracketing interval.
    initial_right: Return value of value_and_gradients_function at the right end
      point of the current bracketing interval.
    f_lim: real `Tensor` of shape [n]. The function value threshold for
      the approximate Wolfe conditions to be checked for each batch member.

  Returns:
    A namedtuple containing the following fields:
      iteration: An int32 scalar `Tensor`. The number of iterations performed.
        Bounded above by `max_iterations` parameter.
      stopped: A boolean scalar `Tensor`. True if the Func algorithm
        terminated.
      failed: A scalar boolean tensor. Indicates whether the objective function
        failed to produce a finite value.
      num_evals: A scalar int32 tensor. The number of value and gradients
        function evaluations.
      left: Return value of value_and_gradients_function at the left end
        point of the bracketing interval found.
      right: Return value of value_and_gradients_function at the right end
        point of the bracketing interval found.
  """
  arg_4 = ~is_finite(arg_1, arg_2)
  arg_5 = (arg_2.df < 0) & (arg_2.f > arg_3)
  arg_6 = _IntermediateResult(
      iteration=tf.convert_to_tensor(value=0),
      stopped=arg_4 | ~arg_5,
      arg_4=arg_4,
      num_evals=tf.convert_to_tensor(value=0),
      left=arg_1,
      right=arg_2)
  return _Func(arg_0, arg_6, arg_3)


def Func(arg_0, arg_1, arg_2):
  """Actual implementation of bisect given initial_args in a _BracketResult."""
  def _loop_cond(arg_3):
    # TODO(b/112524024): Also take into account max_iterations.
    return ~tf.reduce_all(input_tensor=arg_3.stopped)

  def _loop_body(arg_3):
    """Narrow down interval to satisfy opposite slope conditions."""
    arg_4 = arg_0((arg_3.left.x + arg_3.right.x) / 2)

    # Fail if function values at mid point are no longer finite; or left/right
    # points are so close to it that we can't distinguish them any more.
    arg_5 = (arg_3.failed | ~is_finite(arg_4) |
              tf.equal(arg_4.x, arg_3.left.x) | tf.equal(arg_4.x, arg_3.right.x))

    # If mid point has a negative slope and the function value at that point is
    # small enough, we can use it as a new left end point to narrow down the
    # interval. If mid point has a positive slope, then we have found a suitable
    # right end point to bracket a minima within opposite slopes. Otherwise, the
    # mid point has a negative slope but the function value at that point is too
    # high to work as left end point, we are in the same situation in which we
    # started the loop so we just update the right end point and continue.
    arg_6 = ~(arg_3.stopped | arg_5)
    arg_7 = (arg_4.df < 0) & (arg_4.f <= arg_2)
    arg_8 = val_where(arg_6 & arg_7, arg_4, arg_3.left)
    arg_9 = val_where(arg_6 & ~arg_7, arg_4, arg_3.right)

    # We're done when the right end point has a positive slope.
    arg_10 = arg_3.stopped | arg_5 | (arg_9.df >= 0)

    return [_IntermediateResult(
        iteration=arg_3.iteration,
        arg_10=arg_10,
        arg_5=arg_5,
        num_evals=arg_3.num_evals + 1,
        arg_8=arg_8,
        arg_9=arg_9)]

  # The interval needs updating if the right end point has a negative slope and
  # the value of the function at that point is too high. It is not a valid left
  # end point but along with the current left end point, it encloses another
  # minima. The loop above tries to narrow the interval so that it satisfies the
  # opposite slope conditions.
  return tf.while_loop(
      cond=_loop_cond, body=_loop_body, loop_vars=[arg_1])[0]


def Func(arg_0, arg_1=None):
  """Checks if the supplied values are finite.

  Args:
    val_1: A namedtuple instance with the function value and derivative,
      as returned e.g. by value_and_gradients_function evaluations.
    val_2: (Optional) A namedtuple instance with the function value and
      derivative, as returned e.g. by value_and_gradients_function evaluations.

  Returns:
    Func: Scalar boolean `Tensor` indicating whether the function value
      and the derivative in `val_1` (and optionally in `val_2`) are all finite.
  """
  arg_2 = tf.math.Func(arg_0.f) & tf.math.Func(arg_0.df)
  if arg_1 is not None:
    return arg_2 & tf.math.Func(arg_1.f) & tf.math.Func(
        arg_1.df)
  return arg_2