def bisect(value_and_gradients_function,
           initial_left,
           initial_right,
           f_lim):
  """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 bisect 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.
  """
  failed = ~is_finite(initial_left, initial_right)
  needs_bisect = (initial_right.df < 0) & (initial_right.f > f_lim)
  bisect_args = _IntermediateResult(
      iteration=tf.convert_to_tensor(value=0),
      stopped=failed | ~needs_bisect,
      failed=failed,
      num_evals=tf.convert_to_tensor(value=0),
      left=initial_left,
      right=initial_right)
  return _bisect(value_and_gradients_function, bisect_args, f_lim)

def bisect ( value_and_gradients_function , initial_left , initial_right , f_lim ) : failed = ~ is_finite ( initial_left , initial_right ) needs_bisect = ( initial_right . df < 0 ) & ( initial_right . f > f_lim ) bisect_args = _IntermediateResult ( iteration = tf . convert_to_tensor ( value = 0 ) , stopped = failed | ~ needs_bisect , failed = failed , num_evals = tf . convert_to_tensor ( value = 0 ) , left = initial_left , right = initial_right ) return _bisect ( value_and_gradients_function , bisect_args , f_lim )

Bisects an interval and updates to satisfy opposite slope conditions .


def _bisect(value_and_gradients_function, initial_args, f_lim):
  """Actual implementation of bisect given initial_args in a _BracketResult."""
  def _loop_cond(curr):
    # TODO(b/112524024): Also take into account max_iterations.
    return ~tf.reduce_all(input_tensor=curr.stopped)

  def _loop_body(curr):
    """Narrow down interval to satisfy opposite slope conditions."""
    mid = value_and_gradients_function((curr.left.x + curr.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.
    failed = (curr.failed | ~is_finite(mid) |
              tf.equal(mid.x, curr.left.x) | tf.equal(mid.x, curr.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.
    to_update = ~(curr.stopped | failed)
    update_left = (mid.df < 0) & (mid.f <= f_lim)
    left = val_where(to_update & update_left, mid, curr.left)
    right = val_where(to_update & ~update_left, mid, curr.right)

    # We're done when the right end point has a positive slope.
    stopped = curr.stopped | failed | (right.df >= 0)

    return [_IntermediateResult(
        iteration=curr.iteration,
        stopped=stopped,
        failed=failed,
        num_evals=curr.num_evals + 1,
        left=left,
        right=right)]

  # 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=[initial_args])[0]

Actual implementation of bisect given initial_args in a _BracketResult .


def is_finite(val_1, val_2=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:
    is_finite: Scalar boolean `Tensor` indicating whether the function value
      and the derivative in `val_1` (and optionally in `val_2`) are all finite.
  """
  val_1_finite = tf.math.is_finite(val_1.f) & tf.math.is_finite(val_1.df)
  if val_2 is not None:
    return val_1_finite & tf.math.is_finite(val_2.f) & tf.math.is_finite(
        val_2.df)
  return val_1_finite

def is_finite ( val_1 , val_2 = None ) : val_1_finite = tf . math . is_finite ( val_1 . f ) & tf . math . is_finite ( val_1 . df ) if val_2 is not None : return val_1_finite & tf . math . is_finite ( val_2 . f ) & tf . math . is_finite ( val_2 . df ) return val_1_finite

Checks if the supplied values are finite .


def _satisfies_wolfe(val_0,
                     val_c,
                     f_lim,
                     sufficient_decrease_param,
                     curvature_param):
  """Checks whether the Wolfe or approx Wolfe conditions are satisfied.

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

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

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

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

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

  Args:
    val_0: A namedtuple, as returned by value_and_gradients_function
      evaluated at 0.
    val_c: A namedtuple, as returned by value_and_gradients_function
      evaluated at the point to be tested.
    f_lim: Scalar `Tensor` of real dtype. The function value threshold for
      the approximate Wolfe conditions to be checked.
    sufficient_decrease_param: Positive scalar `Tensor` of real dtype.
      Bounded above by the curvature param. Corresponds to 'delta' in the
      terminology of [Hager and Zhang (2006)][2].
    curvature_param: Positive scalar `Tensor` of real dtype. Bounded above
      by `1.`. Corresponds to 'sigma' in the terminology of
      [Hager Zhang (2005)][1].

  Returns:
    is_satisfied: A scalar boolean `Tensor` which is True if either the
      Wolfe or approximate Wolfe conditions are satisfied.
  """
  exact_wolfe_suff_dec = (sufficient_decrease_param * val_0.df >=
                          (val_c.f - val_0.f) / val_c.x)
  wolfe_curvature = val_c.df >= curvature_param * val_0.df
  exact_wolfe = exact_wolfe_suff_dec & wolfe_curvature
  approx_wolfe_applies = val_c.f <= f_lim
  approx_wolfe_suff_dec = ((2 * sufficient_decrease_param - 1) * val_0.df
                           >= val_c.df)
  approx_wolfe = approx_wolfe_applies & approx_wolfe_suff_dec & wolfe_curvature
  is_satisfied = exact_wolfe | approx_wolfe
  return is_satisfied

Checks whether the Wolfe or approx Wolfe conditions are satisfied .


def _secant(val_a, val_b):
  """Returns the secant interpolation for the minimum.

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

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

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

def _secant ( val_a , val_b ) : return ( val_a . x * val_b . df - val_b . x * val_a . df ) / ( val_b . df - val_a . df )

Returns the secant interpolation for the minimum .


def make_simple_step_size_update_policy(num_adaptation_steps,
                                        target_rate=0.75,
                                        decrement_multiplier=0.01,
                                        increment_multiplier=0.01,
                                        step_counter=None):
  """Create a function implementing a step-size update policy.

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

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

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

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

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

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

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

    if kernel_results is None:
      if mcmc_util.is_list_like(step_size_var):
        return [tf.identity(ss) for ss in step_size_var]
      return tf.identity(step_size_var)
    log_n = tf.math.log(
        tf.cast(
            tf.size(input=kernel_results.log_accept_ratio),
            kernel_results.log_accept_ratio.dtype))
    log_mean_accept_ratio = tf.reduce_logsumexp(
        input_tensor=tf.minimum(kernel_results.log_accept_ratio, 0.)) - log_n
    adjustment = tf.where(
        log_mean_accept_ratio < tf.cast(
            tf.math.log(target_rate), log_mean_accept_ratio.dtype),
        -decrement_multiplier / (1. + decrement_multiplier),
        increment_multiplier)

    def build_assign_op():
      if mcmc_util.is_list_like(step_size_var):
        return [
            ss.assign_add(ss * tf.cast(adjustment, ss.dtype))
            for ss in step_size_var
        ]
      return step_size_var.assign_add(
          step_size_var * tf.cast(adjustment, step_size_var.dtype))

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

  return step_size_simple_update_fn

Create a function implementing a step - size update policy .


def _leapfrog_integrator_one_step(
    target_log_prob_fn,
    independent_chain_ndims,
    step_sizes,
    current_momentum_parts,
    current_state_parts,
    current_target_log_prob,
    current_target_log_prob_grad_parts,
    state_gradients_are_stopped=False,
    name=None):
  """Applies `num_leapfrog_steps` of the leapfrog integrator.

  Assumes a simple quadratic kinetic energy function: `0.5 ||momentum||**2`.

  #### Examples:

  ##### Simple quadratic potential.

  ```python
  import matplotlib.pyplot as plt
  %matplotlib inline
  import numpy as np
  import tensorflow as tf
  from tensorflow_probability.python.mcmc.hmc import _leapfrog_integrator_one_step  # pylint: disable=line-too-long
  tfd = tfp.distributions

  dims = 10
  num_iter = int(1e3)
  dtype = np.float32

  position = tf.placeholder(np.float32)
  momentum = tf.placeholder(np.float32)

  target_log_prob_fn = tfd.MultivariateNormalDiag(
      loc=tf.zeros(dims, dtype)).log_prob

  def _leapfrog_one_step(*args):
    # Closure representing computation done during each leapfrog step.
    return _leapfrog_integrator_one_step(
        target_log_prob_fn=target_log_prob_fn,
        independent_chain_ndims=0,
        step_sizes=[0.1],
        current_momentum_parts=args[0],
        current_state_parts=args[1],
        current_target_log_prob=args[2],
        current_target_log_prob_grad_parts=args[3])

  # Do leapfrog integration.
  [
      [next_momentum],
      [next_position],
      next_target_log_prob,
      next_target_log_prob_grad_parts,
  ] = tf.while_loop(
      cond=lambda *args: True,
      body=_leapfrog_one_step,
      loop_vars=[
        [momentum],
        [position],
        target_log_prob_fn(position),
        tf.gradients(target_log_prob_fn(position), position),
      ],
      maximum_iterations=3)

  momentum_ = np.random.randn(dims).astype(dtype)
  position_ = np.random.randn(dims).astype(dtype)
  positions = np.zeros([num_iter, dims], dtype)

  with tf.Session() as sess:
    for i in xrange(num_iter):
      position_, momentum_ = sess.run(
          [next_momentum, next_position],
          feed_dict={position: position_, momentum: momentum_})
      positions[i] = position_

  plt.plot(positions[:, 0]);  # Sinusoidal.
  ```

  Args:
    target_log_prob_fn: Python callable which takes an argument like
      `*current_state_parts` and returns its (possibly unnormalized) log-density
      under the target distribution.
    independent_chain_ndims: Scalar `int` `Tensor` representing the number of
      leftmost `Tensor` dimensions which index independent chains.
    step_sizes: Python `list` of `Tensor`s representing the step size for the
      leapfrog integrator. Must broadcast with the shape of
      `current_state_parts`.  Larger step sizes lead to faster progress, but
      too-large step sizes make rejection exponentially more likely. When
      possible, it's often helpful to match per-variable step sizes to the
      standard deviations of the target distribution in each variable.
    current_momentum_parts: Tensor containing the value(s) of the momentum
      variable(s) to update.
    current_state_parts: Python `list` of `Tensor`s representing the current
      state(s) of the Markov chain(s). The first `independent_chain_ndims` of
      the `Tensor`(s) index different chains.
    current_target_log_prob: `Tensor` representing the value of
      `target_log_prob_fn(*current_state_parts)`. The only reason to specify
      this argument is to reduce TF graph size.
    current_target_log_prob_grad_parts: Python list of `Tensor`s representing
      gradient of `target_log_prob_fn(*current_state_parts`) wrt
      `current_state_parts`. Must have same shape as `current_state_parts`. The
      only reason to specify this argument is to reduce TF graph size.
    state_gradients_are_stopped: Python `bool` indicating that the proposed new
      state be run through `tf.stop_gradient`. This is particularly useful when
      combining optimization over samples from the HMC chain.
      Default value: `False` (i.e., do not apply `stop_gradient`).
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., 'hmc_leapfrog_integrator').

  Returns:
    proposed_momentum_parts: Updated value of the momentum.
    proposed_state_parts: Tensor or Python list of `Tensor`s representing the
      state(s) of the Markov chain(s) at each result step. Has same shape as
      input `current_state_parts`.
    proposed_target_log_prob: `Tensor` representing the value of
      `target_log_prob_fn` at `next_state`.
    proposed_target_log_prob_grad_parts: Gradient of `proposed_target_log_prob`
      wrt `next_state`.

  Raises:
    ValueError: if `len(momentum_parts) != len(state_parts)`.
    ValueError: if `len(state_parts) != len(step_sizes)`.
    ValueError: if `len(state_parts) != len(grads_target_log_prob)`.
    TypeError: if `not target_log_prob.dtype.is_floating`.
  """
  # Note on per-variable step sizes:
  #
  # Using per-variable step sizes is equivalent to using the same step
  # size for all variables and adding a diagonal mass matrix in the
  # kinetic energy term of the Hamiltonian being integrated. This is
  # hinted at by Neal (2011) but not derived in detail there.
  #
  # Let x and v be position and momentum variables respectively.
  # Let g(x) be the gradient of `target_log_prob_fn(x)`.
  # Let S be a diagonal matrix of per-variable step sizes.
  # Let the Hamiltonian H(x, v) = -target_log_prob_fn(x) + 0.5 * ||v||**2.
  #
  # Using per-variable step sizes gives the updates
  # v'  = v  + 0.5 * matmul(S, g(x))
  # x'' = x  + matmul(S, v')
  # v'' = v' + 0.5 * matmul(S, g(x''))
  #
  # Let u = matmul(inv(S), v).
  # Multiplying v by inv(S) in the updates above gives the transformed dynamics
  # u'  = matmul(inv(S), v')  = matmul(inv(S), v) + 0.5 * g(x)
  #                           = u + 0.5 * g(x)
  # x'' = x + matmul(S, v') = x + matmul(S**2, u')
  # u'' = matmul(inv(S), v'') = matmul(inv(S), v') + 0.5 * g(x'')
  #                           = u' + 0.5 * g(x'')
  #
  # These are exactly the leapfrog updates for the Hamiltonian
  # H'(x, u) = -target_log_prob_fn(x) + 0.5 * u^T S**2 u
  #          = -target_log_prob_fn(x) + 0.5 * ||v||**2 = H(x, v).
  #
  # To summarize:
  #
  # * Using per-variable step sizes implicitly simulates the dynamics
  #   of the Hamiltonian H' (which are energy-conserving in H'). We
  #   keep track of v instead of u, but the underlying dynamics are
  #   the same if we transform back.
  # * The value of the Hamiltonian H'(x, u) is the same as the value
  #   of the original Hamiltonian H(x, v) after we transform back from
  #   u to v.
  # * Sampling v ~ N(0, I) is equivalent to sampling u ~ N(0, S**-2).
  #
  # So using per-variable step sizes in HMC will give results that are
  # exactly identical to explicitly using a diagonal mass matrix.

  with tf.compat.v1.name_scope(name, 'hmc_leapfrog_integrator_one_step', [
      independent_chain_ndims, step_sizes, current_momentum_parts,
      current_state_parts, current_target_log_prob,
      current_target_log_prob_grad_parts
  ]):

    # Step 1: Update momentum.
    proposed_momentum_parts = [
        v + 0.5 * tf.cast(eps, v.dtype) * g
        for v, eps, g
        in zip(current_momentum_parts,
               step_sizes,
               current_target_log_prob_grad_parts)]

    # Step 2: Update state.
    proposed_state_parts = [
        x + tf.cast(eps, v.dtype) * v
        for x, eps, v
        in zip(current_state_parts,
               step_sizes,
               proposed_momentum_parts)]

    if state_gradients_are_stopped:
      proposed_state_parts = [tf.stop_gradient(x) for x in proposed_state_parts]

    # Step 3a: Re-evaluate target-log-prob (and grad) at proposed state.
    [
        proposed_target_log_prob,
        proposed_target_log_prob_grad_parts,
    ] = mcmc_util.maybe_call_fn_and_grads(
        target_log_prob_fn,
        proposed_state_parts)

    if not proposed_target_log_prob.dtype.is_floating:
      raise TypeError('`target_log_prob_fn` must produce a `Tensor` '
                      'with `float` `dtype`.')

    if any(g is None for g in proposed_target_log_prob_grad_parts):
      raise ValueError(
          'Encountered `None` gradient. Does your target `target_log_prob_fn` '
          'access all `tf.Variable`s via `tf.get_variable`?\n'
          '  current_state_parts: {}\n'
          '  proposed_state_parts: {}\n'
          '  proposed_target_log_prob_grad_parts: {}'.format(
              current_state_parts,
              proposed_state_parts,
              proposed_target_log_prob_grad_parts))

    # Step 3b: Update momentum (again).
    proposed_momentum_parts = [
        v + 0.5 * tf.cast(eps, v.dtype) * g
        for v, eps, g
        in zip(proposed_momentum_parts,
               step_sizes,
               proposed_target_log_prob_grad_parts)]

    return [
        proposed_momentum_parts,
        proposed_state_parts,
        proposed_target_log_prob,
        proposed_target_log_prob_grad_parts,
    ]

Applies num_leapfrog_steps of the leapfrog integrator .


def _compute_log_acceptance_correction(current_momentums,
                                       proposed_momentums,
                                       independent_chain_ndims,
                                       name=None):
  """Helper to `kernel` which computes the log acceptance-correction.

  A sufficient but not necessary condition for the existence of a stationary
  distribution, `p(x)`, is "detailed balance", i.e.:

  ```none
  p(x'|x) p(x) = p(x|x') p(x')
  ```

  In the Metropolis-Hastings algorithm, a state is proposed according to
  `g(x'|x)` and accepted according to `a(x'|x)`, hence
  `p(x'|x) = g(x'|x) a(x'|x)`.

  Inserting this into the detailed balance equation implies:

  ```none
      g(x'|x) a(x'|x) p(x) = g(x|x') a(x|x') p(x')
  ==> a(x'|x) / a(x|x') = p(x') / p(x) [g(x|x') / g(x'|x)]    (*)
  ```

  One definition of `a(x'|x)` which satisfies (*) is:

  ```none
  a(x'|x) = min(1, p(x') / p(x) [g(x|x') / g(x'|x)])
  ```

  (To see that this satisfies (*), notice that under this definition only at
  most one `a(x'|x)` and `a(x|x') can be other than one.)

  We call the bracketed term the "acceptance correction".

  In the case of UncalibratedHMC, the log acceptance-correction is not the log
  proposal-ratio. UncalibratedHMC augments the state-space with momentum, z.
  Assuming a standard Gaussian distribution for momentums, the chain eventually
  converges to:

  ```none
  p([x, z]) propto= target_prob(x) exp(-0.5 z**2)
  ```

  Relating this back to Metropolis-Hastings parlance, for HMC we have:

  ```none
  p([x, z]) propto= target_prob(x) exp(-0.5 z**2)
  g([x, z] | [x', z']) = g([x', z'] | [x, z])
  ```

  In other words, the MH bracketed term is `1`. However, because we desire to
  use a general MH framework, we can place the momentum probability ratio inside
  the metropolis-correction factor thus getting an acceptance probability:

  ```none
                       target_prob(x')
  accept_prob(x'|x) = -----------------  [exp(-0.5 z**2) / exp(-0.5 z'**2)]
                       target_prob(x)
  ```

  (Note: we actually need to handle the kinetic energy change at each leapfrog
  step, but this is the idea.)

  Args:
    current_momentums: `Tensor` representing the value(s) of the current
      momentum(s) of the state (parts).
    proposed_momentums: `Tensor` representing the value(s) of the proposed
      momentum(s) of the state (parts).
    independent_chain_ndims: Scalar `int` `Tensor` representing the number of
      leftmost `Tensor` dimensions which index independent chains.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., 'compute_log_acceptance_correction').

  Returns:
    log_acceptance_correction: `Tensor` representing the `log`
      acceptance-correction.  (See docstring for mathematical definition.)
  """
  with tf.compat.v1.name_scope(
      name, 'compute_log_acceptance_correction',
      [independent_chain_ndims, current_momentums, proposed_momentums]):
    log_current_kinetic, log_proposed_kinetic = [], []
    for current_momentum, proposed_momentum in zip(
        current_momentums, proposed_momentums):
      axis = tf.range(independent_chain_ndims, tf.rank(current_momentum))
      log_current_kinetic.append(_log_sum_sq(current_momentum, axis))
      log_proposed_kinetic.append(_log_sum_sq(proposed_momentum, axis))
    current_kinetic = 0.5 * tf.exp(
        tf.reduce_logsumexp(
            input_tensor=tf.stack(log_current_kinetic, axis=-1), axis=-1))
    proposed_kinetic = 0.5 * tf.exp(
        tf.reduce_logsumexp(
            input_tensor=tf.stack(log_proposed_kinetic, axis=-1), axis=-1))
    return mcmc_util.safe_sum([current_kinetic, -proposed_kinetic])

def _compute_log_acceptance_correction ( current_momentums , proposed_momentums , independent_chain_ndims , name = None ) : with tf . compat . v1 . name_scope ( name , 'compute_log_acceptance_correction' , [ independent_chain_ndims , current_momentums , proposed_momentums ] ) : log_current_kinetic , log_proposed_kinetic = [ ] , [ ] for current_momentum , proposed_momentum in zip ( current_momentums , proposed_momentums ) : axis = tf . range ( independent_chain_ndims , tf . rank ( current_momentum ) ) log_current_kinetic . append ( _log_sum_sq ( current_momentum , axis ) ) log_proposed_kinetic . append ( _log_sum_sq ( proposed_momentum , axis ) ) current_kinetic = 0.5 * tf . exp ( tf . reduce_logsumexp ( input_tensor = tf . stack ( log_current_kinetic , axis = - 1 ) , axis = - 1 ) ) proposed_kinetic = 0.5 * tf . exp ( tf . reduce_logsumexp ( input_tensor = tf . stack ( log_proposed_kinetic , axis = - 1 ) , axis = - 1 ) ) return mcmc_util . safe_sum ( [ current_kinetic , - proposed_kinetic ] )

Helper to kernel which computes the log acceptance - correction .


def one_step(self, current_state, previous_kernel_results):
    """Runs one iteration of Hamiltonian Monte Carlo.

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

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

    Raises:
      ValueError: if there isn't one `step_size` or a list with same length as
        `current_state`.
    """
    previous_step_size_assign = (
        [] if self.step_size_update_fn is None
        else (previous_kernel_results.extra.step_size_assign
              if mcmc_util.is_list_like(
                  previous_kernel_results.extra.step_size_assign)
              else [previous_kernel_results.extra.step_size_assign]))

    with tf.control_dependencies(previous_step_size_assign):
      next_state, kernel_results = self._impl.one_step(
          current_state, previous_kernel_results)
      if self.step_size_update_fn is not None:
        step_size_assign = self.step_size_update_fn(  # pylint: disable=not-callable
            self.step_size, kernel_results)
        kernel_results = kernel_results._replace(
            extra=HamiltonianMonteCarloExtraKernelResults(
                step_size_assign=step_size_assign))
      return next_state, kernel_results

def
one_step
(
self
,
current_state
,
previous_kernel_results
)
:
previous_step_size_assign
=
(
[
]
if
self
.
step_size_update_fn
is
else
(
previous_kernel_results
.
extra
.
step_size_assign
if
mcmc_util
.
is_list_like
(
previous_kernel_results
.
extra
.
step_size_assign
)
else
[
previous_kernel_results
.
extra
.
step_size_assign
]
)
)
with
tf
.
control_dependencies
(
previous_step_size_assign
)
:
next_state
,
kernel_results
=
self
.
_impl
.
one_step
(
current_state
,
previous_kernel_results
)
if
self
.
step_size_update_fn
is
not
:
step_size_assign
=
self
.
step_size_update_fn
(
# pylint: disable=not-callable
self
.
step_size
,
kernel_results
)
kernel_results
=
kernel_results
.
_replace
(
extra
=
HamiltonianMonteCarloExtraKernelResults
(
step_size_assign
=
step_size_assign
)
)
return
next_state
,
kernel_results

Runs one iteration of Hamiltonian Monte Carlo .


def bootstrap_results(self, init_state):
    """Creates initial `previous_kernel_results` using a supplied `state`."""
    kernel_results = self._impl.bootstrap_results(init_state)
    if self.step_size_update_fn is not None:
      step_size_assign = self.step_size_update_fn(self.step_size, None)  # pylint: disable=not-callable
      kernel_results = kernel_results._replace(
          extra=HamiltonianMonteCarloExtraKernelResults(
              step_size_assign=step_size_assign))
    return kernel_results

def
bootstrap_results
(
self
,
init_state
)
:
kernel_results
=
self
.
_impl
.
bootstrap_results
(
init_state
)
if
self
.
step_size_update_fn
is
not
:
step_size_assign
=
self
.
step_size_update_fn
(
self
.
step_size
,
)
# pylint: disable=not-callable
kernel_results
=
kernel_results
.
_replace
(
extra
=
HamiltonianMonteCarloExtraKernelResults
(
step_size_assign
=
step_size_assign
)
)
return
kernel_results

Creates initial previous_kernel_results using a supplied state .


def bayesian_resnet(input_shape,
                    num_classes=10,
                    kernel_posterior_scale_mean=-9.0,
                    kernel_posterior_scale_stddev=0.1,
                    kernel_posterior_scale_constraint=0.2):
  """Constructs a ResNet18 model.

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

  Returns:
    tf.keras.Model.
  """

  filters = [64, 128, 256, 512]
  kernels = [3, 3, 3, 3]
  strides = [1, 2, 2, 2]

  def _untransformed_scale_constraint(t):
    return tf.clip_by_value(t, -1000,
                            tf.math.log(kernel_posterior_scale_constraint))

  kernel_posterior_fn = tfp.layers.default_mean_field_normal_fn(
      untransformed_scale_initializer=tf.compat.v1.initializers.random_normal(
          mean=kernel_posterior_scale_mean,
          stddev=kernel_posterior_scale_stddev),
      untransformed_scale_constraint=_untransformed_scale_constraint)

  image = tf.keras.layers.Input(shape=input_shape, dtype='float32')
  x = tfp.layers.Convolution2DFlipout(
      64,
      3,
      strides=1,
      padding='same',
      kernel_posterior_fn=kernel_posterior_fn)(image)

  for i in range(len(kernels)):
    x = _resnet_block(
        x,
        filters[i],
        kernels[i],
        strides[i],
        kernel_posterior_fn)

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

  x = tfp.layers.DenseFlipout(
      num_classes,
      kernel_posterior_fn=kernel_posterior_fn)(x)

  model = tf.keras.Model(inputs=image, outputs=x, name='resnet18')
  return model

def bayesian_resnet ( input_shape , num_classes = 10 , kernel_posterior_scale_mean = - 9.0 , kernel_posterior_scale_stddev = 0.1 , kernel_posterior_scale_constraint = 0.2 ) : filters = [ 64 , 128 , 256 , 512 ] kernels = [ 3 , 3 , 3 , 3 ] strides = [ 1 , 2 , 2 , 2 ] def _untransformed_scale_constraint ( t ) : return tf . clip_by_value ( t , - 1000 , tf . math . log ( kernel_posterior_scale_constraint ) ) kernel_posterior_fn = tfp . layers . default_mean_field_normal_fn ( untransformed_scale_initializer = tf . compat . v1 . initializers . random_normal ( mean = kernel_posterior_scale_mean , stddev = kernel_posterior_scale_stddev ) , untransformed_scale_constraint = _untransformed_scale_constraint ) image = tf . keras . layers . Input ( shape = input_shape , dtype = 'float32' ) x = tfp . layers . Convolution2DFlipout ( 64 , 3 , strides = 1 , padding = 'same' , kernel_posterior_fn = kernel_posterior_fn ) ( image ) for i in range ( len ( kernels ) ) : x = _resnet_block ( x , filters [ i ] , kernels [ i ] , strides [ i ] , kernel_posterior_fn ) x = tf . keras . layers . BatchNormalization ( ) ( x ) x = tf . keras . layers . Activation ( 'relu' ) ( x ) x = tf . keras . layers . AveragePooling2D ( 4 , 1 ) ( x ) x = tf . keras . layers . Flatten ( ) ( x ) x = tfp . layers . DenseFlipout ( num_classes , kernel_posterior_fn = kernel_posterior_fn ) ( x ) model = tf . keras . Model ( inputs = image , outputs = x , name = 'resnet18' ) return model


def _resnet_block(x, filters, kernel, stride, kernel_posterior_fn):
  """Network block for ResNet."""
  x = tf.keras.layers.BatchNormalization()(x)
  x = tf.keras.layers.Activation('relu')(x)

  if stride != 1 or filters != x.shape[1]:
    shortcut = _projection_shortcut(x, filters, stride, kernel_posterior_fn)
  else:
    shortcut = x

  x = tfp.layers.Convolution2DFlipout(
      filters,
      kernel,
      strides=stride,
      padding='same',
      kernel_posterior_fn=kernel_posterior_fn)(x)
  x = tf.keras.layers.BatchNormalization()(x)
  x = tf.keras.layers.Activation('relu')(x)

  x = tfp.layers.Convolution2DFlipout(
      filters,
      kernel,
      strides=1,
      padding='same',
      kernel_posterior_fn=kernel_posterior_fn)(x)
  x = tf.keras.layers.add([x, shortcut])
  return x

def _resnet_block ( x , filters , kernel , stride , kernel_posterior_fn ) : x = tf . keras . layers . BatchNormalization ( ) ( x ) x = tf . keras . layers . Activation ( 'relu' ) ( x ) if stride != 1 or filters != x . shape [ 1 ] : shortcut = _projection_shortcut ( x , filters , stride , kernel_posterior_fn ) else : shortcut = x x = tfp . layers . Convolution2DFlipout ( filters , kernel , strides = stride , padding = 'same' , kernel_posterior_fn = kernel_posterior_fn ) ( x ) x = tf . keras . layers . BatchNormalization ( ) ( x ) x = tf . keras . layers . Activation ( 'relu' ) ( x ) x = tfp . layers . Convolution2DFlipout ( filters , kernel , strides = 1 , padding = 'same' , kernel_posterior_fn = kernel_posterior_fn ) ( x ) x = tf . keras . layers . add ( [ x , shortcut ] ) return x


def make_encoder(activation, num_topics, layer_sizes):
  """Create the encoder function.

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

  Returns:
    encoder: A `callable` mapping a bag-of-words `Tensor` to a
      `tfd.Distribution` instance over topics.
  """
  encoder_net = tf.keras.Sequential()
  for num_hidden_units in layer_sizes:
    encoder_net.add(
        tf.keras.layers.Dense(
            num_hidden_units,
            activation=activation,
            kernel_initializer=tf.compat.v1.glorot_normal_initializer()))
  encoder_net.add(
      tf.keras.layers.Dense(
          num_topics,
          activation=tf.nn.softplus,
          kernel_initializer=tf.compat.v1.glorot_normal_initializer()))

  def encoder(bag_of_words):
    net = _clip_dirichlet_parameters(encoder_net(bag_of_words))
    return tfd.Dirichlet(concentration=net,
                         name="topics_posterior")

  return encoder

def make_encoder ( activation , num_topics , layer_sizes ) : encoder_net = tf . keras . Sequential ( ) for num_hidden_units in layer_sizes : encoder_net . add ( tf . keras . layers . Dense ( num_hidden_units , activation = activation , kernel_initializer = tf . compat . v1 . glorot_normal_initializer ( ) ) ) encoder_net . add ( tf . keras . layers . Dense ( num_topics , activation = tf . nn . softplus , kernel_initializer = tf . compat . v1 . glorot_normal_initializer ( ) ) ) def encoder ( bag_of_words ) : net = _clip_dirichlet_parameters ( encoder_net ( bag_of_words ) ) return tfd . Dirichlet ( concentration = net , name = "topics_posterior" ) return encoder


def make_decoder(num_topics, num_words):
  """Create the decoder function.

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

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

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

  return decoder, topics_words


def make_prior(num_topics, initial_value):
  """Create the prior distribution.

  Args:
    num_topics: Number of topics.
    initial_value: The starting value for the prior parameters.

  Returns:
    prior: A `callable` that returns a `tf.distribution.Distribution`
        instance, the prior distribution.
    prior_variables: A `list` of `Variable` objects, the trainable parameters
        of the prior.
  """
  def _softplus_inverse(x):
    return np.log(np.expm1(x))

  logit_concentration = tf.compat.v1.get_variable(
      "logit_concentration",
      shape=[1, num_topics],
      initializer=tf.compat.v1.initializers.constant(
          _softplus_inverse(initial_value)))
  concentration = _clip_dirichlet_parameters(
      tf.nn.softplus(logit_concentration))

  def prior():
    return tfd.Dirichlet(concentration=concentration,
                         name="topics_prior")

  prior_variables = [logit_concentration]

  return prior, prior_variables

def make_prior ( num_topics , initial_value ) : def _softplus_inverse ( x ) : return np . log ( np . expm1 ( x ) ) logit_concentration = tf . compat . v1 . get_variable ( "logit_concentration" , shape = [ 1 , num_topics ] , initializer = tf . compat . v1 . initializers . constant ( _softplus_inverse ( initial_value ) ) ) concentration = _clip_dirichlet_parameters ( tf . nn . softplus ( logit_concentration ) ) def prior ( ) : return tfd . Dirichlet ( concentration = concentration , name = "topics_prior" ) prior_variables = [ logit_concentration ] return prior , prior_variables


def sample_chain(
    num_results,
    current_state,
    previous_kernel_results=None,
    kernel=None,
    num_burnin_steps=0,
    num_steps_between_results=0,
    trace_fn=lambda current_state, kernel_results: kernel_results,
    return_final_kernel_results=False,
    parallel_iterations=10,
    name=None,
):
  """Implements Markov chain Monte Carlo via repeated `TransitionKernel` steps.

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

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

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

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

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

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

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

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

  #### Examples

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

  I.e.,

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

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

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

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

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

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

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

  I.e.,

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

  where `F` denotes factors.

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

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

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

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

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

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

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

  # Sample from Hamiltonian Monte Carlo Markov Chain.

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

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

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

  ##### Custom tracing functions.

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

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

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

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

  def trace_everything(states, previous_kernel_results):
    return previous_kernel_results

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

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

  #### References

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

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

    def _trace_scan_fn(state_and_results, num_steps):
      next_state, current_kernel_results = mcmc_util.smart_for_loop(
          loop_num_iter=num_steps,
          body_fn=kernel.one_step,
          initial_loop_vars=list(state_and_results),
          parallel_iterations=parallel_iterations)
      return next_state, current_kernel_results

    (_, final_kernel_results), (all_states, trace) = mcmc_util.trace_scan(
        loop_fn=_trace_scan_fn,
        initial_state=(current_state, previous_kernel_results),
        elems=tf.one_hot(
            indices=0,
            depth=num_results,
            on_value=1 + num_burnin_steps,
            off_value=1 + num_steps_between_results,
            dtype=tf.int32),
        # pylint: disable=g-long-lambda
        trace_fn=lambda state_and_results: (state_and_results[0],
                                            trace_fn(*state_and_results)),
        # pylint: enable=g-long-lambda
        parallel_iterations=parallel_iterations)

    if return_final_kernel_results:
      return CheckpointableStatesAndTrace(
          all_states=all_states,
          trace=trace,
          final_kernel_results=final_kernel_results)
    else:
      if no_trace:
        return all_states
      else:
        return StatesAndTrace(all_states=all_states, trace=trace)

Implements Markov chain Monte Carlo via repeated TransitionKernel steps .


def deep_exponential_family(data_size, feature_size, units, shape):
  """A multi-layered topic model over a documents-by-terms matrix."""
  w2 = ed.Gamma(0.1, 0.3, sample_shape=[units[2], units[1]], name="w2")
  w1 = ed.Gamma(0.1, 0.3, sample_shape=[units[1], units[0]], name="w1")
  w0 = ed.Gamma(0.1, 0.3, sample_shape=[units[0], feature_size], name="w0")

  z2 = ed.Gamma(0.1, 0.1, sample_shape=[data_size, units[2]], name="z2")
  z1 = ed.Gamma(shape, shape / tf.matmul(z2, w2), name="z1")
  z0 = ed.Gamma(shape, shape / tf.matmul(z1, w1), name="z0")
  x = ed.Poisson(tf.matmul(z0, w0), name="x")
  return x

def deep_exponential_family ( data_size , feature_size , units , shape ) : w2 = ed . Gamma ( 0.1 , 0.3 , sample_shape = [ units [ 2 ] , units [ 1 ] ] , name = "w2" ) w1 = ed . Gamma ( 0.1 , 0.3 , sample_shape = [ units [ 1 ] , units [ 0 ] ] , name = "w1" ) w0 = ed . Gamma ( 0.1 , 0.3 , sample_shape = [ units [ 0 ] , feature_size ] , name = "w0" ) z2 = ed . Gamma ( 0.1 , 0.1 , sample_shape = [ data_size , units [ 2 ] ] , name = "z2" ) z1 = ed . Gamma ( shape , shape / tf . matmul ( z2 , w2 ) , name = "z1" ) z0 = ed . Gamma ( shape , shape / tf . matmul ( z1 , w1 ) , name = "z0" ) x = ed . Poisson ( tf . matmul ( z0 , w0 ) , name = "x" ) return x

A multi - layered topic model over a documents - by - terms matrix .


def trainable_positive_deterministic(shape, min_loc=1e-3, name=None):
  """Learnable Deterministic distribution over positive reals."""
  with tf.compat.v1.variable_scope(
      None, default_name="trainable_positive_deterministic"):
    unconstrained_loc = tf.compat.v1.get_variable("unconstrained_loc", shape)
    loc = tf.maximum(tf.nn.softplus(unconstrained_loc), min_loc)
    rv = ed.Deterministic(loc=loc, name=name)
    return rv

def trainable_positive_deterministic ( shape , min_loc = 1e-3 , name = None ) : with tf . compat . v1 . variable_scope ( None , default_name = "trainable_positive_deterministic" ) : unconstrained_loc = tf . compat . v1 . get_variable ( "unconstrained_loc" , shape ) loc = tf . maximum ( tf . nn . softplus ( unconstrained_loc ) , min_loc ) rv = ed . Deterministic ( loc = loc , name = name ) return rv

Learnable Deterministic distribution over positive reals .


def trainable_gamma(shape, min_concentration=1e-3, min_scale=1e-5, name=None):
  """Learnable Gamma via concentration and scale parameterization."""
  with tf.compat.v1.variable_scope(None, default_name="trainable_gamma"):
    unconstrained_concentration = tf.compat.v1.get_variable(
        "unconstrained_concentration",
        shape,
        initializer=tf.compat.v1.initializers.random_normal(
            mean=0.5, stddev=0.1))
    unconstrained_scale = tf.compat.v1.get_variable(
        "unconstrained_scale",
        shape,
        initializer=tf.compat.v1.initializers.random_normal(stddev=0.1))
    concentration = tf.maximum(tf.nn.softplus(unconstrained_concentration),
                               min_concentration)
    rate = tf.maximum(1. / tf.nn.softplus(unconstrained_scale), 1. / min_scale)
    rv = ed.Gamma(concentration=concentration, rate=rate, name=name)
    return rv

def trainable_gamma ( shape , min_concentration = 1e-3 , min_scale = 1e-5 , name = None ) : with tf . compat . v1 . variable_scope ( None , default_name = "trainable_gamma" ) : unconstrained_concentration = tf . compat . v1 . get_variable ( "unconstrained_concentration" , shape , initializer = tf . compat . v1 . initializers . random_normal ( mean = 0.5 , stddev = 0.1 ) ) unconstrained_scale = tf . compat . v1 . get_variable ( "unconstrained_scale" , shape , initializer = tf . compat . v1 . initializers . random_normal ( stddev = 0.1 ) ) concentration = tf . maximum ( tf . nn . softplus ( unconstrained_concentration ) , min_concentration ) rate = tf . maximum ( 1. / tf . nn . softplus ( unconstrained_scale ) , 1. / min_scale ) rv = ed . Gamma ( concentration = concentration , rate = rate , name = name ) return rv

Learnable Gamma via concentration and scale parameterization .


def load_nips2011_papers(path):
  """Loads NIPS 2011 conference papers.

  The NIPS 1987-2015 data set is in the form of a 11,463 x 5,812 matrix of
  per-paper word counts, containing 11,463 words and 5,811 NIPS conference
  papers (Perrone et al., 2016). We subset to papers in 2011 and words appearing
  in at least two documents and having a total word count of at least 10.

  Built from the Observations Python package.

  Args:
    path: str.
      Path to directory which either stores file or otherwise file will
      be downloaded and extracted there. Filename is `NIPS_1987-2015.csv`.

  Returns:
    bag_of_words: np.ndarray of shape [num_documents, num_words]. Each element
      denotes the number of occurrences of a specific word in a specific
      document.
    words: List of strings, denoting the words for `bag_of_words`'s columns.
  """
  path = os.path.expanduser(path)
  filename = "NIPS_1987-2015.csv"
  filepath = os.path.join(path, filename)
  if not os.path.exists(filepath):
    url = ("https://archive.ics.uci.edu/ml/machine-learning-databases/"
           "00371/NIPS_1987-2015.csv")
    if not tf.io.gfile.exists(path):
      tf.io.gfile.makedirs(path)
    print("Downloading %s to %s" % (url, filepath))
    urllib.request.urlretrieve(url, filepath)

  with open(filepath) as f:
    iterator = csv.reader(f)
    documents = next(iterator)[1:]
    words = []
    x_train = []
    for row in iterator:
      words.append(row[0])
      x_train.append(row[1:])

  x_train = np.array(x_train, dtype=np.int)

  # Subset to documents in 2011 and words appearing in at least two documents
  # and have a total word count of at least 10.
  doc_idx = [i for i, document in enumerate(documents)
             if document.startswith("2011")]
  documents = [documents[doc] for doc in doc_idx]
  x_train = x_train[:, doc_idx]
  word_idx = np.logical_and(np.sum(x_train != 0, 1) >= 2,
                            np.sum(x_train, 1) >= 10)
  words = [word for word, idx in zip(words, word_idx) if idx]
  bag_of_words = x_train[word_idx, :].T
  return bag_of_words, words


def _init_params(self, amplitude, length_scale, validate_args):
    """Shared init logic for `amplitude` and `length_scale` params.

    Args:
      amplitude: `Tensor` (or convertible) or `None` to convert, validate.
      length_scale: `Tensor` (or convertible) or `None` to convert, validate.
      validate_args: If `True`, parameters are checked for validity despite
        possibly degrading runtime performance

    Returns:
      dtype: The common `DType` of the parameters.
    """
    dtype = util.maybe_get_common_dtype(
        [amplitude, length_scale])
    if amplitude is not None:
      amplitude = tf.convert_to_tensor(
          value=amplitude, name='amplitude', dtype=dtype)
    self._amplitude = _validate_arg_if_not_none(
        amplitude, tf.compat.v1.assert_positive, validate_args)
    if length_scale is not None:
      length_scale = tf.convert_to_tensor(
          value=length_scale, name='length_scale', dtype=dtype)
    self._length_scale = _validate_arg_if_not_none(
        length_scale, tf.compat.v1.assert_positive, validate_args)
    return dtype

def _init_params ( self , amplitude , length_scale , validate_args ) : dtype = util . maybe_get_common_dtype ( [ amplitude , length_scale ] ) if amplitude is not None : amplitude = tf . convert_to_tensor ( value = amplitude , name = 'amplitude' , dtype = dtype ) self . _amplitude = _validate_arg_if_not_none ( amplitude , tf . compat . v1 . assert_positive , validate_args ) if length_scale is not None : length_scale = tf . convert_to_tensor ( value = length_scale , name = 'length_scale' , dtype = dtype ) self . _length_scale = _validate_arg_if_not_none ( length_scale , tf . compat . v1 . assert_positive , validate_args ) return dtype

Shared init logic for amplitude and length_scale params .


def _registered_kl(type_a, type_b):
  """Get the KL function registered for classes a and b."""
  hierarchy_a = tf_inspect.getmro(type_a)
  hierarchy_b = tf_inspect.getmro(type_b)
  dist_to_children = None
  kl_fn = None
  for mro_to_a, parent_a in enumerate(hierarchy_a):
    for mro_to_b, parent_b in enumerate(hierarchy_b):
      candidate_dist = mro_to_a + mro_to_b
      candidate_kl_fn = _DIVERGENCES.get((parent_a, parent_b), None)
      if not kl_fn or (candidate_kl_fn and candidate_dist < dist_to_children):
        dist_to_children = candidate_dist
        kl_fn = candidate_kl_fn
  return kl_fn

Get the KL function registered for classes a and b .


def read_image(filepath):
  """Returns an image tensor."""
  im_bytes = tf.io.read_file(filepath)
  im = tf.image.decode_image(im_bytes, channels=CHANNELS)
  im = tf.image.convert_image_dtype(im, tf.float32)
  return im

def read_image ( filepath ) : im_bytes = tf . io . read_file ( filepath ) im = tf . image . decode_image ( im_bytes , channels = CHANNELS ) im = tf . image . convert_image_dtype ( im , tf . float32 ) return im


def download_sprites():
  """Downloads the sprites data and returns the saved filepath."""
  filepath = os.path.join(FLAGS.data_dir, DATA_SPRITES_DIR)
  if not tf.io.gfile.exists(filepath):
    if not tf.io.gfile.exists(FLAGS.data_dir):
      tf.io.gfile.makedirs(FLAGS.data_dir)
    zip_name = "{}.zip".format(filepath)
    urllib.request.urlretrieve(DATA_SPRITES_URL, zip_name)
    with zipfile.ZipFile(zip_name, "r") as zip_file:
      zip_file.extractall(FLAGS.data_dir)
    tf.io.gfile.remove(zip_name)
  return filepath

def download_sprites ( ) : filepath = os . path . join ( FLAGS . data_dir , DATA_SPRITES_DIR ) if not tf . io . gfile . exists ( filepath ) : if not tf . io . gfile . exists ( FLAGS . data_dir ) : tf . io . gfile . makedirs ( FLAGS . data_dir ) zip_name = "{}.zip" . format ( filepath ) urllib . request . urlretrieve ( DATA_SPRITES_URL , zip_name ) with zipfile . ZipFile ( zip_name , "r" ) as zip_file : zip_file . extractall ( FLAGS . data_dir ) tf . io . gfile . remove ( zip_name ) return filepath

Downloads the sprites data and returns the saved filepath .


def create_character(skin, hair, top, pants):
  """Creates a character sprite from a set of attribute sprites."""
  dtype = skin.dtype
  hair_mask = tf.cast(hair[..., -1:] <= 0, dtype)
  top_mask = tf.cast(top[..., -1:] <= 0, dtype)
  pants_mask = tf.cast(pants[..., -1:] <= 0, dtype)
  char = (skin * hair_mask) + hair
  char = (char * top_mask) + top
  char = (char * pants_mask) + pants
  return char

def create_character ( skin , hair , top , pants ) : dtype = skin . dtype hair_mask = tf . cast ( hair [ ... , - 1 : ] <= 0 , dtype ) top_mask = tf . cast ( top [ ... , - 1 : ] <= 0 , dtype ) pants_mask = tf . cast ( pants [ ... , - 1 : ] <= 0 , dtype ) char = ( skin * hair_mask ) + hair char = ( char * top_mask ) + top char = ( char * pants_mask ) + pants return char

Creates a character sprite from a set of attribute sprites .


def create_seq(character, action_metadata, direction, length=8, start=0):
  """Creates a sequence.

  Args:
    character: A character sprite tensor.
    action_metadata: An action metadata tuple.
    direction: An integer representing the direction, i.e., the row
      offset within each action group corresponding to a particular
      direction.
    length: Desired length of the sequence. If this is longer than
      the number of available frames, it will roll over to the
      beginning.
    start: Index of possible frames at which to start the sequence.

  Returns:
    A sequence tensor.
  """
  sprite_start = (action_metadata[0]+direction) * FRAME_SIZE
  sprite_end = (action_metadata[0]+direction+1) * FRAME_SIZE
  sprite_line = character[sprite_start:sprite_end, ...]

  # Extract 64x64 patches that are side-by-side in the sprite, and limit
  # to the actual number of frames for the given action.
  frames = tf.stack(tf.split(sprite_line, 13, axis=1))  # 13 is a hack
  frames = frames[0:action_metadata[1]]

  # Extract a slice of the desired length.
  # NOTE: Length could be longer than the number of frames, so tile as needed.
  frames = tf.roll(frames, shift=-start, axis=0)
  frames = tf.tile(frames, [2, 1, 1, 1])  # 2 is a hack
  frames = frames[:length]
  frames = tf.cast(frames, dtype=tf.float32)
  frames.set_shape([length, FRAME_SIZE, FRAME_SIZE, CHANNELS])
  return frames


def create_random_seq(character, action_metadata, direction, length=8):
  """Creates a random sequence."""
  start = tf.random.uniform([], maxval=action_metadata[1], dtype=tf.int32)
  return create_seq(character, action_metadata, direction, length, start)

def create_random_seq ( character , action_metadata , direction , length = 8 ) : start = tf . random . uniform ( [ ] , maxval = action_metadata [ 1 ] , dtype = tf . int32 ) return create_seq ( character , action_metadata , direction , length , start )


def create_sprites_dataset(characters, actions, directions, channels=3,
                           length=8, shuffle=False, fake_data=False):
  """Creates a tf.data pipeline for the sprites dataset.

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

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

  action_names = [action.name for action in actions]
  action_metadata = [(action.start_row, action.frames) for action in actions]

  direction_rows = [direction.row_offset for direction in directions]

  chars = tf.data.Dataset.from_tensor_slices(characters)
  act_names = tf.data.Dataset.from_tensor_slices(action_names).repeat()
  acts_metadata = tf.data.Dataset.from_tensor_slices(action_metadata).repeat()
  dir_rows = tf.data.Dataset.from_tensor_slices(direction_rows).repeat()

  if shuffle:
    chars = chars.shuffle(len(characters))

  dataset = tf.data.Dataset.zip((chars, act_names, acts_metadata, dir_rows))

  skin_table = tf.contrib.lookup.index_table_from_tensor(sorted(SKIN_COLORS))
  hair_table = tf.contrib.lookup.index_table_from_tensor(sorted(HAIRSTYLES))
  top_table = tf.contrib.lookup.index_table_from_tensor(sorted(TOPS))
  pants_table = tf.contrib.lookup.index_table_from_tensor(sorted(PANTS))
  action_table = tf.contrib.lookup.index_table_from_tensor(sorted(action_names))

  def process_example(attrs, act_name, act_metadata, dir_row_offset):
    """Processes a dataset row."""
    skin_name = attrs[0]
    hair_name = attrs[1]
    top_name = attrs[2]
    pants_name = attrs[3]

    if fake_data:
      char = dummy_image
    else:
      skin = read_image(basedir + os.sep + skin_name)
      hair = read_image(basedir + os.sep + hair_name)
      top = read_image(basedir + os.sep + top_name)
      pants = read_image(basedir + os.sep + pants_name)
      char = create_character(skin, hair, top, pants)

    if shuffle:
      seq = create_random_seq(char, act_metadata, dir_row_offset, length)
    else:
      seq = create_seq(char, act_metadata, dir_row_offset, length)
    seq = seq[..., :channels]  # limit output channels

    skin_idx = skin_table.lookup(skin_name)
    hair_idx = hair_table.lookup(hair_name)
    top_idx = top_table.lookup(top_name)
    pants_idx = pants_table.lookup(pants_name)
    act_idx = action_table.lookup(act_name)

    return (seq, skin_idx, hair_idx, top_idx, pants_idx, act_idx,
            skin_name, hair_name, top_name, pants_name, act_name)

  dataset = dataset.map(process_example)
  return dataset

def
create_sprites_dataset
(
characters
,
actions
,
directions
,
channels
=
3
,
length
=
8
,
shuffle
=
False
,
fake_data
=
False
)
:
if
fake_data
:
dummy_image
=
tf
.
random
.
normal
(
[
HEIGHT
,
WIDTH
,
CHANNELS
]
)
else
:
basedir
=
download_sprites
(
)
action_names
=
[
action
.
name
for
action
in
actions
]
action_metadata
=
[
(
action
.
start_row
,
action
.
frames
)
for
action
in
actions
]
direction_rows
=
[
direction
.
row_offset
for
direction
in
directions
]
chars
=
tf
.
data
.
Dataset
.
from_tensor_slices
(
characters
)
act_names
=
tf
.
data
.
Dataset
.
from_tensor_slices
(
action_names
)
.
repeat
(
)
acts_metadata
=
tf
.
data
.
Dataset
.
from_tensor_slices
(
action_metadata
)
.
repeat
(
)
dir_rows
=
tf
.
data
.
Dataset
.
from_tensor_slices
(
direction_rows
)
.
repeat
(
)
if
shuffle
:
chars
=
chars
.
shuffle
(
len
(
characters
)
)
dataset
=
tf
.
data
.
Dataset
.
zip
(
(
chars
,
act_names
,
acts_metadata
,
dir_rows
)
)
skin_table
=
tf
.
contrib
.
lookup
.
index_table_from_tensor
(
sorted
(
SKIN_COLORS
)
)
hair_table
=
tf
.
contrib
.
lookup
.
index_table_from_tensor
(
sorted
(
HAIRSTYLES
)
)
top_table
=
tf
.
contrib
.
lookup
.
index_table_from_tensor
(
sorted
(
TOPS
)
)
pants_table
=
tf
.
contrib
.
lookup
.
index_table_from_tensor
(
sorted
(
PANTS
)
)
action_table
=
tf
.
contrib
.
lookup
.
index_table_from_tensor
(
sorted
(
action_names
)
)
def
process_example
(
attrs
,
act_name
,
act_metadata
,
dir_row_offset
)
:
"""Processes a dataset row."""
skin_name
=
attrs
[
0
]
hair_name
=
attrs
[
1
]
top_name
=
attrs
[
2
]
pants_name
=
attrs
[
3
]
if
fake_data
:
char
=
dummy_image
else
:
skin
=
read_image
(
basedir
+
os
.
sep
+
skin_name
)
hair
=
read_image
(
basedir
+
os
.
sep
+
hair_name
)
top
=
read_image
(
basedir
+
os
.
sep
+
top_name
)
pants
=
read_image
(
basedir
+
os
.
sep
+
pants_name
)
char
=
create_character
(
skin
,
hair
,
top
,
pants
)
if
shuffle
:
seq
=
create_random_seq
(
char
,
act_metadata
,
dir_row_offset
,
length
)
else
:
seq
=
create_seq
(
char
,
act_metadata
,
dir_row_offset
,
length
)
seq
=
seq
[
...
,
:
channels
]
# limit output channels
skin_idx
=
skin_table
.
lookup
(
skin_name
)
hair_idx
=
hair_table
.
lookup
(
hair_name
)
top_idx
=
top_table
.
lookup
(
top_name
)
pants_idx
=
pants_table
.
lookup
(
pants_name
)
act_idx
=
action_table
.
lookup
(
act_name
)
return
(
seq
,
skin_idx
,
hair_idx
,
top_idx
,
pants_idx
,
act_idx
,
skin_name
,
hair_name
,
top_name
,
pants_name
,
act_name
)
dataset
=
dataset
.
map
(
process_example
)
return
dataset

Creates a tf . data pipeline for the sprites dataset .


def _maybe_validate_distributions(distributions, dtype_override, validate_args):
  """Checks that `distributions` satisfies all assumptions."""
  assertions = []

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

  if dtype_override is None:
    dts = [
        dtype_util.base_dtype(d.dtype)
        for d in distributions
        if d.dtype is not None
    ]
    if dts[1:] != dts[:-1]:
      raise TypeError('Distributions must have same dtype; found: {}.'.format(
          set(dtype_util.name(dt) for dt in dts)))

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

  batch_shapes = [d.batch_shape for d in distributions]
  if all(tensorshape_util.is_fully_defined(b) for b in batch_shapes):
    if batch_shapes[1:] != batch_shapes[:-1]:
      raise ValueError('Distributions must have the same `batch_shape`; '
                       'found: {}.'.format(batch_shapes))
  elif validate_args:
    batch_shapes = [
        tensorshape_util.as_list(d.batch_shape)  # pylint: disable=g-complex-comprehension
        if tensorshape_util.is_fully_defined(d.batch_shape) else
        d.batch_shape_tensor() for d in distributions
    ]
    assertions.extend(
        assert_util.assert_equal(  # pylint: disable=g-complex-comprehension
            b1, b2,
            message='Distribution `batch_shape`s must be identical.')
        for b1, b2 in zip(batch_shapes[1:], batch_shapes[:-1]))

  return assertions

Checks that distributions satisfies all assumptions .


def _flatten_summand_list(kernels):
  """Flatten a list of kernels which may contain _SumKernel instances.

  Args:
    kernels: Python list of `PositiveSemidefiniteKernel` instances

  Returns:
    Python list containing the elements of kernels, with any _SumKernel
    instances replaced by their `kernels` property contents.
  """
  flattened = []
  for k in kernels:
    if isinstance(k, _SumKernel):
      flattened += k.kernels
    else:
      flattened.append(k)
  return flattened

def _flatten_summand_list ( kernels ) : flattened = [ ] for k in kernels : if isinstance ( k , _SumKernel ) : flattened += k . kernels else : flattened . append ( k ) return flattened

Flatten a list of kernels which may contain _SumKernel instances .


def _flatten_multiplicand_list(kernels):
  """Flatten a list of kernels which may contain _ProductKernel instances.

  Args:
    kernels: Python list of `PositiveSemidefiniteKernel` instances

  Returns:
    Python list containing the elements of kernels, with any _ProductKernel
    instances replaced by their `kernels` property contents.
  """
  flattened = []
  for k in kernels:
    if isinstance(k, _ProductKernel):
      flattened += k.kernels
    else:
      flattened.append(k)
  return flattened

def _flatten_multiplicand_list ( kernels ) : flattened = [ ] for k in kernels : if isinstance ( k , _ProductKernel ) : flattened += k . kernels else : flattened . append ( k ) return flattened

Flatten a list of kernels which may contain _ProductKernel instances .


def build_fake_data():
  """Build fake CIFAR10-style data for unit testing."""
  num_examples = 10
  x_train = np.random.rand(num_examples, *IMAGE_SHAPE).astype(np.float32)
  y_train = np.random.permutation(np.arange(num_examples)).astype(np.int32)
  x_test = np.random.rand(num_examples, *IMAGE_SHAPE).astype(np.float32)
  y_test = np.random.permutation(np.arange(num_examples)).astype(np.int32)
  return (x_train, y_train), (x_test, y_test)

def build_fake_data ( ) : num_examples = 10 x_train = np . random . rand ( num_examples , * IMAGE_SHAPE ) . astype ( np . float32 ) y_train = np . random . permutation ( np . arange ( num_examples ) ) . astype ( np . int32 ) x_test = np . random . rand ( num_examples , * IMAGE_SHAPE ) . astype ( np . float32 ) y_test = np . random . permutation ( np . arange ( num_examples ) ) . astype ( np . int32 ) return ( x_train , y_train ) , ( x_test , y_test )

Build fake CIFAR10 - style data for unit testing .


def count_integers(arr,
                   weights=None,
                   minlength=None,
                   maxlength=None,
                   axis=None,
                   dtype=tf.int32,
                   name=None):
  """Counts the number of occurrences of each value in an integer array `arr`.

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

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

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

  Returns:
    A vector with the same dtype as `weights` or the given `dtype`. The bin
    values.
  """
  with tf.compat.v1.name_scope(
      name, 'count_integers', values=[arr, weights, minlength, maxlength,
                                      axis]):
    if axis is None:
      return tf.math.bincount(
          arr,
          weights=weights,
          minlength=minlength,
          maxlength=maxlength,
          dtype=dtype)

    arr = tf.convert_to_tensor(value=arr, dtype=tf.int32, name='arr')
    arr_ndims = _get_static_ndims(arr, expect_static=True)

    axis = _make_static_axis_non_negative_list(axis, arr_ndims)

    # ~axis from docstring.  Dims in arr that are not in axis.
    not_axis = sorted(set(range(arr_ndims)).difference(axis))

    # If we're reducing over everything, just use standard bincount.
    if not not_axis:
      return tf.math.bincount(
          arr,
          weights=weights,
          minlength=minlength,
          maxlength=maxlength,
          dtype=dtype)

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

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

      def one_bincount(arr_slice):
        return tf.math.bincount(
            arr_slice,
            weights=None,
            minlength=minlength,
            maxlength=maxlength,
            dtype=dtype)

      flat_counts = tf.map_fn(one_bincount, elems=flat_arr, dtype=dtype)
    else:
      weights = tf.convert_to_tensor(value=weights, name='weights')
      _get_static_ndims(weights, expect_static=True, expect_ndims=arr_ndims)
      flat_weights = _move_dims_to_flat_end(
          weights, not_axis, arr_ndims, right_end=False)

      def one_bincount(arr_and_weights_slices):
        arr_slice, weights_slice = arr_and_weights_slices
        return tf.math.bincount(
            arr_slice,
            weights=weights_slice,
            minlength=minlength,
            maxlength=maxlength,
            dtype=dtype)

      flat_counts = tf.map_fn(
          one_bincount, elems=[flat_arr, flat_weights], dtype=weights.dtype)

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

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

    # not_axis_shape = arr.shape[~axis]
    not_axis_shape = tf.gather(tf.shape(input=arr), indices=not_axis)

    # The first index of flat_counts_t indexes bins 0,..,K-1, the rest are ~axis
    out_shape = tf.concat([[-1], not_axis_shape], axis=0)

    return tf.reshape(flat_counts_t, out_shape)

Counts the number of occurrences of each value in an integer array arr .


def find_bins(x,
              edges,
              extend_lower_interval=False,
              extend_upper_interval=False,
              dtype=None,
              name=None):
  """Bin values into discrete intervals.

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

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

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

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

  #### Examples

  Cut a `1-D` array

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

  Cut `x` into its deciles

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

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

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

    flattening_x = edges.shape.ndims == 1 and x.shape.ndims > 1

    if flattening_x:
      x_orig_shape = tf.shape(input=x)
      x = tf.reshape(x, [-1])

    if dtype is None:
      dtype = in_type
    dtype = tf.as_dtype(dtype)

    # Move first dims into the rightmost.
    x_permed = distribution_util.rotate_transpose(x, shift=-1)
    edges_permed = distribution_util.rotate_transpose(edges, shift=-1)

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

    # In above example, we want [0, 0, 1, 1], so correct this here.
    bins = tf.clip_by_value(almost_output - 1, tf.cast(0, dtype),
                            tf.cast(tf.shape(input=edges)[0] - 2, dtype))

    if not extend_lower_interval:
      low_fill = np.nan if dtype.is_floating else -1
      bins = tf.where(x < tf.expand_dims(edges[0], 0),
                      tf.fill(tf.shape(input=x), tf.cast(low_fill, dtype)),
                      bins)

    if not extend_upper_interval:
      up_fill = np.nan if dtype.is_floating else tf.shape(input=edges)[0] - 1
      bins = tf.where(x > tf.expand_dims(edges[-1], 0),
                      tf.fill(tf.shape(input=x), tf.cast(up_fill, dtype)), bins)

    if flattening_x:
      bins = tf.reshape(bins, x_orig_shape)

    return bins


def histogram(x,
              edges,
              axis=None,
              extend_lower_interval=False,
              extend_upper_interval=False,
              dtype=None,
              name=None):
  """Count how often `x` falls in intervals defined by `edges`.

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

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

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

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

  #### Examples

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

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

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

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

  """
  with tf.compat.v1.name_scope(name, 'histogram', values=[x, edges, axis]):

    # Tensor conversions.
    in_dtype = dtype_util.common_dtype([x, edges], preferred_dtype=tf.float32)

    x = tf.convert_to_tensor(value=x, name='x', dtype=in_dtype)
    edges = tf.convert_to_tensor(value=edges, name='edges', dtype=in_dtype)

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

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

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

Count how often x falls in intervals defined by edges .


def quantiles(x,
              num_quantiles,
              axis=None,
              interpolation=None,
              keep_dims=False,
              validate_args=False,
              name=None):
  """Compute quantiles of `x` along `axis`.

  The quantiles of a distribution are cut points dividing the range into
  intervals with equal probabilities.

  Given a vector `x` of samples, this function estimates the cut points by
  returning `num_quantiles + 1` cut points, `(c0, ..., cn)`, such that, roughly
  speaking, equal number of sample points lie in the `num_quantiles` intervals
  `[c0, c1), [c1, c2), ..., [c_{n-1}, cn]`.  That is,

  * About `1 / n` fraction of the data lies in `[c_{k-1}, c_k)`, `k = 1, ..., n`
  * About `k / n` fraction of the data lies below `c_k`.
  * `c0` is the sample minimum and `cn` is the maximum.

  The exact number of data points in each interval depends on the size of
  `x` (e.g. whether the size is divisible by `n`) and the `interpolation` kwarg.

  Args:
    x:  Numeric `N-D` `Tensor` with `N > 0`.  If `axis` is not `None`,
      `x` must have statically known number of dimensions.
    num_quantiles:  Scalar `integer` `Tensor`.  The number of intervals the
      returned `num_quantiles + 1` cut points divide the range into.
    axis:  Optional `0-D` or `1-D` integer `Tensor` with constant values. The
      axis that index independent samples over which to return the desired
      percentile.  If `None` (the default), treat every dimension as a sample
      dimension, returning a scalar.
    interpolation : {'nearest', 'linear', 'lower', 'higher', 'midpoint'}.
      Default value: 'nearest'.  This specifies the interpolation method to
      use when the fractions `k / n` lie between two data points `i < j`:
        * linear: i + (j - i) * fraction, where fraction is the fractional part
          of the index surrounded by i and j.
        * lower: `i`.
        * higher: `j`.
        * nearest: `i` or `j`, whichever is nearest.
        * midpoint: (i + j) / 2. `linear` and `midpoint` interpolation do not
          work with integer dtypes.
    keep_dims:  Python `bool`. If `True`, the last dimension is kept with size 1
      If `False`, the last dimension is removed from the output shape.
    validate_args:  Whether to add runtime checks of argument validity. If
      False, and arguments are incorrect, correct behavior is not guaranteed.
    name:  A Python string name to give this `Op`.  Default is 'percentile'

  Returns:
    cut_points:  A `rank(x) + 1 - len(axis)` dimensional `Tensor` with same
    `dtype` as `x` and shape `[num_quantiles + 1, ...]` where the trailing shape
    is that of `x` without the dimensions in `axis` (unless `keep_dims is True`)

  Raises:
    ValueError:  If argument 'interpolation' is not an allowed type.
    ValueError:  If interpolation type not compatible with `dtype`.

  #### Examples

  ```python
  # Get quartiles of x with various interpolation choices.
  x = [0.,  1.,   2.,   3.,   4.,   5.,   6.,   7.,   8.,   9.,  10.]

  tfp.stats.quantiles(x, num_quantiles=4, interpolation='nearest')
  ==> [  0.,   2.,   5.,   8.,  10.]

  tfp.stats.quantiles(x, num_quantiles=4, interpolation='linear')
  ==> [  0. ,   2.5,   5. ,   7.5,  10. ]

  tfp.stats.quantiles(x, num_quantiles=4, interpolation='lower')
  ==> [  0.,   2.,   5.,   7.,  10.]

  # Get deciles of columns of an R x C data set.
  data = load_my_columnar_data(...)
  tfp.stats.quantiles(data, num_quantiles=10)
  ==> Shape [11, C] Tensor
  ```

  """
  with tf.compat.v1.name_scope(
      name, 'quantiles', values=[x, num_quantiles, axis]):
    x = tf.convert_to_tensor(value=x, name='x')
    return percentile(
        x,
        q=tf.linspace(
            # percentile casts q to float64 before using it...so may as well use
            # float64 here. Note that  using x.dtype won't work with linspace
            # if x is integral type (which is anothe motivation for hard-coding
            # float64).
            tf.convert_to_tensor(value=0, dtype=tf.float64),
            tf.convert_to_tensor(value=100, dtype=tf.float64),
            num=num_quantiles + 1),
        axis=axis,
        interpolation=interpolation,
        keep_dims=keep_dims,
        validate_args=validate_args,
        preserve_gradients=False)


def _get_static_ndims(x,
                      expect_static=False,
                      expect_ndims=None,
                      expect_ndims_no_more_than=None,
                      expect_ndims_at_least=None):
  """Get static number of dimensions and assert that some expectations are met.

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

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

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

  Returns:
    ndims:  A Python integer.

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

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

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

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

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

  return ndims

Get static number of dimensions and assert that some expectations are met .


def _insert_back_keep_dims(x, axis):
  """Insert the dims in `axis` back as singletons after being removed.

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

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

def _insert_back_keep_dims ( x , axis ) : for i in sorted ( axis ) : x = tf . expand_dims ( x , axis = i ) return x

Insert the dims in axis back as singletons after being removed .


def _make_static_axis_non_negative_list(axis, ndims):
  """Convert possibly negatively indexed axis to non-negative list of ints.

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

  Returns:
    A list of non-negative Python integers.

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

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

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

  return list(int(dim) for dim in axis)

Convert possibly negatively indexed axis to non - negative list of ints .