def _get_permutations(num_results, dims, seed=None):
  """Uniform iid sample from the space of permutations.

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

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

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

def _get_permutations ( num_results , dims , seed = None ) : sample_range = tf . range ( num_results ) stream = distributions . SeedStream ( seed , salt = 'MCMCSampleHaltonSequence3' ) def generate_one ( d ) : seed = stream ( ) fn = lambda _ : tf . random . shuffle ( tf . range ( d ) , seed = seed ) return tf . map_fn ( fn , sample_range , parallel_iterations = 1 if seed is not None else 10 ) return tf . concat ( [ generate_one ( d ) for d in tf . unstack ( dims ) ] , axis = - 1 )

Uniform iid sample from the space of permutations .


def _get_indices(num_results, sequence_indices, dtype, name=None):
  """Generates starting points for the Halton sequence procedure.

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

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

  Returns:
    indices: `Tensor` of dtype `dtype` and shape = `[n, 1, 1]`.
  """
  with tf.compat.v1.name_scope(name, '_get_indices',
                               [num_results, sequence_indices]):
    if sequence_indices is None:
      num_results = tf.cast(num_results, dtype=dtype)
      sequence_indices = tf.range(num_results, dtype=dtype)
    else:
      sequence_indices = tf.cast(sequence_indices, dtype)

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

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

def _get_indices ( num_results , sequence_indices , dtype , name = None ) : with tf . compat . v1 . name_scope ( name , '_get_indices' , [ num_results , sequence_indices ] ) : if sequence_indices is None : num_results = tf . cast ( num_results , dtype = dtype ) sequence_indices = tf . range ( num_results , dtype = dtype ) else : sequence_indices = tf . cast ( sequence_indices , dtype ) indices = sequence_indices + 1 return tf . reshape ( indices , [ - 1 , 1 , 1 ] )

Generates starting points for the Halton sequence procedure .


def _base_expansion_size(num, bases):
  """Computes the number of terms in the place value expansion.

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

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

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

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

def _base_expansion_size ( num , bases ) : return tf . floor ( tf . math . log ( num ) / tf . math . log ( bases ) ) + 1

Computes the number of terms in the place value expansion .


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

def _primes_less_than ( n ) : small_primes = np . array ( ( 2 , 3 , 5 ) ) if n <= 6 : return small_primes [ small_primes < n ] sieve = np . ones ( n // 3 + ( n % 6 == 2 ) , dtype = np . bool ) sieve [ 0 ] = False m = int ( n ** 0.5 ) // 3 + 1 for i in range ( m ) : if not sieve [ i ] : continue k = 3 * i + 1 | 1 sieve [ k ** 2 // 3 : : 2 * k ] = False sieve [ ( k ** 2 + 4 * k - 2 * k * ( i & 1 ) ) // 3 : : 2 * k ] = False return np . r_ [ 2 , 3 , 3 * np . nonzero ( sieve ) [ 0 ] + 1 | 1 ]

Returns sorted array of primes such that 2 < = prime < n .


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

def _machine_eps ( dtype ) : if isinstance ( dtype , tf . DType ) : dtype = dtype . as_numpy_dtype ( ) return np . finfo ( dtype ) . eps

Returns the machine epsilon for the supplied dtype .


def hager_zhang(value_and_gradients_function,
                initial_step_size=None,
                value_at_initial_step=None,
                value_at_zero=None,
                converged=None,
                threshold_use_approximate_wolfe_condition=1e-6,
                shrinkage_param=0.66,
                expansion_param=5.0,
                sufficient_decrease_param=0.1,
                curvature_param=0.9,
                step_size_shrink_param=0.1,
                max_iterations=50,
                name=None):
  """The Hager Zhang line search algorithm.

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

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

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

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

  ### Usage:

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

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

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

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

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

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

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

  Returns:
    results: A namedtuple containing the following attributes.
      converged: Boolean `Tensor` of shape [n]. Whether a point satisfying
        Wolfe/Approx wolfe was found.
      failed: Boolean `Tensor` of shape [n]. Whether line search failed e.g.
        if either the objective function or the gradient are not finite at
        an evaluation point.
      iterations: Scalar int32 `Tensor`. Number of line search iterations made.
      func_evals: Scalar int32 `Tensor`. Number of function evaluations made.
      left: A namedtuple, as returned by value_and_gradients_function,
        of the left end point of the final bracketing interval. Values are
        equal to those of `right` on batch members where converged is True.
        Otherwise, it corresponds to the last interval computed.
      right: A namedtuple, as returned by value_and_gradients_function,
        of the right end point of the final bracketing interval. Values are
        equal to those of `left` on batch members where converged is True.
        Otherwise, it corresponds to the last interval computed.
  """
  with tf.compat.v1.name_scope(name, 'hager_zhang', [
      initial_step_size, value_at_initial_step, value_at_zero, converged,
      threshold_use_approximate_wolfe_condition, shrinkage_param,
      expansion_param, sufficient_decrease_param, curvature_param]):
    val_0, val_initial, f_lim, prepare_evals = _prepare_args(
        value_and_gradients_function,
        initial_step_size,
        value_at_initial_step,
        value_at_zero,
        threshold_use_approximate_wolfe_condition)

    valid_inputs = (hzl.is_finite(val_0) & (val_0.df < 0) &
                    tf.math.is_finite(val_initial.x) & (val_initial.x > 0))

    if converged is None:
      init_converged = tf.zeros_like(valid_inputs)  # i.e. all false.
    else:
      init_converged = tf.convert_to_tensor(value=converged)

    failed = ~init_converged & ~valid_inputs
    active = ~init_converged & valid_inputs

    # Note: _fix_step_size returns immediately if either all inputs are invalid
    # or none of the active ones need fixing.
    fix_step_evals, val_c, fix_failed = _fix_step_size(
        value_and_gradients_function, val_initial, active,
        step_size_shrink_param)

    init_interval = HagerZhangLineSearchResult(
        converged=init_converged,
        failed=failed | fix_failed,
        func_evals=prepare_evals + fix_step_evals,
        iterations=tf.convert_to_tensor(value=0),
        left=val_0,
        right=hzl.val_where(init_converged, val_0, val_c))

    def _apply_bracket_and_search():
      """Bracketing and searching to do for valid inputs."""
      return _bracket_and_search(
          value_and_gradients_function, init_interval, f_lim, max_iterations,
          shrinkage_param, expansion_param, sufficient_decrease_param,
          curvature_param)

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


def _fix_step_size(value_and_gradients_function,
                   val_c_input,
                   active,
                   step_size_shrink_param):
  """Shrinks the input step size until the value and grad become finite."""
  # The maximum iterations permitted are determined as the number of halvings
  # it takes to reduce 1 to 0 in the given dtype.
  iter_max = np.ceil(-np.log2(_machine_eps(val_c_input.x.dtype)))

  def _cond(i, val_c, to_fix):
    del val_c  # Unused.
    return (i < iter_max) & tf.reduce_any(input_tensor=to_fix)

  def _body(i, val_c, to_fix):
    next_c = tf.where(to_fix, val_c.x * step_size_shrink_param, val_c.x)
    next_val_c = value_and_gradients_function(next_c)
    still_to_fix = to_fix & ~hzl.is_finite(next_val_c)
    return (i + 1, next_val_c, still_to_fix)

  to_fix = active & ~hzl.is_finite(val_c_input)
  return tf.while_loop(
      cond=_cond, body=_body, loop_vars=(0, val_c_input, to_fix))

def _fix_step_size ( value_and_gradients_function , val_c_input , active , step_size_shrink_param ) : iter_max = np . ceil ( - np . log2 ( _machine_eps ( val_c_input . x . dtype ) ) ) def _cond ( i , val_c , to_fix ) : del val_c return ( i < iter_max ) & tf . reduce_any ( input_tensor = to_fix ) def _body ( i , val_c , to_fix ) : next_c = tf . where ( to_fix , val_c . x * step_size_shrink_param , val_c . x ) next_val_c = value_and_gradients_function ( next_c ) still_to_fix = to_fix & ~ hzl . is_finite ( next_val_c ) return ( i + 1 , next_val_c , still_to_fix ) to_fix = active & ~ hzl . is_finite ( val_c_input ) return tf . while_loop ( cond = _cond , body = _body , loop_vars = ( 0 , val_c_input , to_fix ) )

Shrinks the input step size until the value and grad become finite .


def _bracket_and_search(
    value_and_gradients_function,
    init_interval,
    f_lim,
    max_iterations,
    shrinkage_param,
    expansion_param,
    sufficient_decrease_param,
    curvature_param):
  """Brackets the minimum and performs a line search.

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

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

  converged = init_interval.converged | _very_close(
      bracket_result.left.x, bracket_result.right.x)

  # We fail if we have not yet converged but already exhausted all iterations.
  exhausted_iterations = ~converged & tf.greater_equal(
      bracket_result.iteration, max_iterations)

  line_search_args = HagerZhangLineSearchResult(
      converged=converged,
      failed=bracket_result.failed | exhausted_iterations,
      iterations=bracket_result.iteration,
      func_evals=bracket_result.num_evals,
      left=bracket_result.left,
      right=bracket_result.right)

  return _line_search_after_bracketing(
      value_and_gradients_function, line_search_args, init_interval.left,
      f_lim, max_iterations, sufficient_decrease_param, curvature_param,
      shrinkage_param)

def _bracket_and_search ( value_and_gradients_function , init_interval , f_lim , max_iterations , shrinkage_param , expansion_param , sufficient_decrease_param , curvature_param ) : bracket_result = hzl . bracket ( value_and_gradients_function , init_interval , f_lim , max_iterations , expansion_param ) converged = init_interval . converged | _very_close ( bracket_result . left . x , bracket_result . right . x ) exhausted_iterations = ~ converged & tf . greater_equal ( bracket_result . iteration , max_iterations ) line_search_args = HagerZhangLineSearchResult ( converged = converged , failed = bracket_result . failed | exhausted_iterations , iterations = bracket_result . iteration , func_evals = bracket_result . num_evals , left = bracket_result . left , right = bracket_result . right ) return _line_search_after_bracketing ( value_and_gradients_function , line_search_args , init_interval . left , f_lim , max_iterations , sufficient_decrease_param , curvature_param , shrinkage_param )

Brackets the minimum and performs a line search .


def _line_search_after_bracketing(
    value_and_gradients_function,
    search_interval,
    val_0,
    f_lim,
    max_iterations,
    sufficient_decrease_param,
    curvature_param,
    shrinkage_param):
  """The main loop of line search after the minimum has been bracketed.

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

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

  def _loop_cond(curr_interval):
    """Loop condition."""
    active = ~(curr_interval.converged | curr_interval.failed)
    return (curr_interval.iterations <
            max_iterations) & tf.reduce_any(input_tensor=active)

  def _loop_body(curr_interval):
    """The loop body."""
    secant2_raw_result = hzl.secant2(
        value_and_gradients_function, val_0, curr_interval, f_lim,
        sufficient_decrease_param, curvature_param)
    secant2_result = HagerZhangLineSearchResult(
        converged=secant2_raw_result.converged,
        failed=secant2_raw_result.failed,
        iterations=curr_interval.iterations + 1,
        func_evals=secant2_raw_result.num_evals,
        left=secant2_raw_result.left,
        right=secant2_raw_result.right)

    should_check_shrinkage = ~(secant2_result.converged | secant2_result.failed)

    def _do_check_shrinkage():
      """Check if interval has shrinked enough."""
      old_width = curr_interval.right.x - curr_interval.left.x
      new_width = secant2_result.right.x - secant2_result.left.x
      sufficient_shrinkage = new_width < old_width * shrinkage_param
      func_is_flat = (
          _very_close(curr_interval.left.f, curr_interval.right.f) &
          _very_close(secant2_result.left.f, secant2_result.right.f))

      new_converged = (
          should_check_shrinkage & sufficient_shrinkage & func_is_flat)
      needs_inner_bisect = should_check_shrinkage & ~sufficient_shrinkage

      inner_bisect_args = secant2_result._replace(
          converged=secant2_result.converged | new_converged)

      def _apply_inner_bisect():
        return _line_search_inner_bisection(
            value_and_gradients_function, inner_bisect_args,
            needs_inner_bisect, f_lim)

      return prefer_static.cond(
          tf.reduce_any(input_tensor=needs_inner_bisect),
          _apply_inner_bisect,
          lambda: inner_bisect_args)

    next_args = prefer_static.cond(
        tf.reduce_any(input_tensor=should_check_shrinkage),
        _do_check_shrinkage,
        lambda: secant2_result)

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

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

def _line_search_after_bracketing ( value_and_gradients_function , search_interval , val_0 , f_lim , max_iterations , sufficient_decrease_param , curvature_param , shrinkage_param ) : def _loop_cond ( curr_interval ) : active = ~ ( curr_interval . converged | curr_interval . failed ) return ( curr_interval . iterations < max_iterations ) & tf . reduce_any ( input_tensor = active ) def _loop_body ( curr_interval ) : secant2_raw_result = hzl . secant2 ( value_and_gradients_function , val_0 , curr_interval , f_lim , sufficient_decrease_param , curvature_param ) secant2_result = HagerZhangLineSearchResult ( converged = secant2_raw_result . converged , failed = secant2_raw_result . failed , iterations = curr_interval . iterations + 1 , func_evals = secant2_raw_result . num_evals , left = secant2_raw_result . left , right = secant2_raw_result . right ) should_check_shrinkage = ~ ( secant2_result . converged | secant2_result . failed ) def _do_check_shrinkage ( ) : old_width = curr_interval . right . x - curr_interval . left . x new_width = secant2_result . right . x - secant2_result . left . x sufficient_shrinkage = new_width < old_width * shrinkage_param func_is_flat = ( _very_close ( curr_interval . left . f , curr_interval . right . f ) & _very_close ( secant2_result . left . f , secant2_result . right . f ) ) new_converged = ( should_check_shrinkage & sufficient_shrinkage & func_is_flat ) needs_inner_bisect = should_check_shrinkage & ~ sufficient_shrinkage inner_bisect_args = secant2_result . _replace ( converged = secant2_result . converged | new_converged ) def _apply_inner_bisect ( ) : return _line_search_inner_bisection ( value_and_gradients_function , inner_bisect_args , needs_inner_bisect , f_lim ) return prefer_static . cond ( tf . reduce_any ( input_tensor = needs_inner_bisect ) , _apply_inner_bisect , lambda : inner_bisect_args ) next_args = prefer_static . cond ( tf . reduce_any ( input_tensor = should_check_shrinkage ) , _do_check_shrinkage , lambda : secant2_result ) interval_shrunk = ( ~ next_args . failed & _very_close ( next_args . left . x , next_args . right . x ) ) return [ next_args . _replace ( converged = next_args . converged | interval_shrunk ) ] return tf . while_loop ( cond = _loop_cond , body = _loop_body , loop_vars = [ search_interval ] , parallel_iterations = 1 ) [ 0 ]

The main loop of line search after the minimum has been bracketed .


def _line_search_inner_bisection(
    value_and_gradients_function,
    search_interval,
    active,
    f_lim):
  """Performs bisection and updates the interval."""
  midpoint = (search_interval.left.x + search_interval.right.x) / 2
  val_mid = value_and_gradients_function(midpoint)
  is_valid_mid = hzl.is_finite(val_mid)

  still_active = active & is_valid_mid
  new_failed = active & ~is_valid_mid
  next_inteval = search_interval._replace(
      failed=search_interval.failed | new_failed,
      func_evals=search_interval.func_evals + 1)

  def _apply_update():
    update_result = hzl.update(
        value_and_gradients_function, next_inteval.left, next_inteval.right,
        val_mid, f_lim, active=still_active)
    return HagerZhangLineSearchResult(
        converged=next_inteval.converged,
        failed=next_inteval.failed | update_result.failed,
        iterations=next_inteval.iterations + update_result.iteration,
        func_evals=next_inteval.func_evals + update_result.num_evals,
        left=update_result.left,
        right=update_result.right)

  return prefer_static.cond(
      tf.reduce_any(input_tensor=still_active),
      _apply_update,
      lambda: next_inteval)

def _line_search_inner_bisection ( value_and_gradients_function , search_interval , active , f_lim ) : midpoint = ( search_interval . left . x + search_interval . right . x ) / 2 val_mid = value_and_gradients_function ( midpoint ) is_valid_mid = hzl . is_finite ( val_mid ) still_active = active & is_valid_mid new_failed = active & ~ is_valid_mid next_inteval = search_interval . _replace ( failed = search_interval . failed | new_failed , func_evals = search_interval . func_evals + 1 ) def _apply_update ( ) : update_result = hzl . update ( value_and_gradients_function , next_inteval . left , next_inteval . right , val_mid , f_lim , active = still_active ) return HagerZhangLineSearchResult ( converged = next_inteval . converged , failed = next_inteval . failed | update_result . failed , iterations = next_inteval . iterations + update_result . iteration , func_evals = next_inteval . func_evals + update_result . num_evals , left = update_result . left , right = update_result . right ) return prefer_static . cond ( tf . reduce_any ( input_tensor = still_active ) , _apply_update , lambda : next_inteval )

Performs bisection and updates the interval .


def _prepare_args(value_and_gradients_function,
                  initial_step_size,
                  val_initial,
                  val_0,
                  approximate_wolfe_threshold):
  """Prepares the arguments for the line search initialization.

  Args:
    value_and_gradients_function: A Python callable that accepts a real scalar
      tensor and returns a namedtuple with the fields 'x', 'f', and 'df' that
      correspond to scalar tensors of real dtype containing the point at which
      the function was evaluated, the value of the function, and its
      derivative at that point. The other namedtuple fields, if present,
      should be tensors or sequences (possibly nested) of tensors.
      In usual optimization application, this function would be generated by
      projecting the multivariate objective function along some specific
      direction. The direction is determined by some other procedure but should
      be a descent direction (i.e. the derivative of the projected univariate
      function must be negative at 0.).
      Alternatively, the function may represent the batching of `n` such line
      functions (e.g. projecting a single multivariate objective function along
      `n` distinct directions at once) accepting n points as input, i.e. a
      tensor of shape [n], and the fields 'x', 'f' and 'df' in the returned
      namedtuple should each be a tensor of shape [n], with the corresponding
      input points, function values, and derivatives at those input points.
    initial_step_size: Scalar positive `Tensor` of real dtype, or a tensor of
      shape [n] in batching mode. The initial value (or values) to try to
      bracket the minimum. Default is `1.` as a float32.
      Note that this point need not necessarily bracket the minimum for the line
      search to work correctly but the supplied value must be greater than 0.
      A good initial value will make the search converge faster.
    val_initial: The full return value of evaluating
      value_and_gradients_function at initial_step_size, i.e. a namedtuple with
      'x', 'f', 'df', if already known by the caller. If not None the value of
      `initial_step_size` will be ignored, otherwise the tuple will be computed
      by evaluating value_and_gradients_function.
    val_0: The full return value of value_and_gradients_function at `0.`, i.e.
      a namedtuple with 'x', 'f', 'df', if already known by the caller. If None
      the tuple will be computed by evaluating value_and_gradients_function.
    approximate_wolfe_threshold: Scalar positive `Tensor` of
      real dtype. Corresponds to the parameter 'epsilon' in
      [Hager and Zhang (2006)][2]. Used to estimate the
      threshold at which the line search switches to approximate Wolfe
      conditions.

  Returns:
    left: A namedtuple, as returned by value_and_gradients_function,
      containing the value and derivative of the function at `0.`.
    val_initial: A namedtuple, as returned by value_and_gradients_function,
      containing the value and derivative of the function at
      `initial_step_size`.
    f_lim: Real `Tensor` of shape [n]. The function value threshold for
      the approximate Wolfe conditions to be checked.
    eval_count: Scalar int32 `Tensor`. The number of target function
      evaluations made by this function.
  """
  eval_count = 0
  if val_initial is None:
    if initial_step_size is not None:
      initial_step_size = tf.convert_to_tensor(value=initial_step_size)
    else:
      initial_step_size = tf.convert_to_tensor(value=1.0, dtype=tf.float32)
    val_initial = value_and_gradients_function(initial_step_size)
    eval_count += 1

  if val_0 is None:
    x_0 = tf.zeros_like(val_initial.x)
    val_0 = value_and_gradients_function(x_0)
    eval_count += 1

  f_lim = val_0.f + (approximate_wolfe_threshold * tf.abs(val_0.f))
  return val_0, val_initial, f_lim, tf.convert_to_tensor(value=eval_count)

Prepares the arguments for the line search initialization .


def _print(pass_through_tensor, values):
  """Wrapper for tf.Print which supports lists and namedtuples for printing."""
  flat_values = []
  for value in values:
    # Checks if it is a namedtuple.
    if hasattr(value, '_fields'):
      for field in value._fields:
        flat_values.extend([field, _to_str(getattr(value, field))])
      continue
    if isinstance(value, (list, tuple)):
      for v in value:
        flat_values.append(_to_str(v))
      continue
    flat_values.append(_to_str(value))
  return tf.compat.v1.Print(pass_through_tensor, flat_values)

def _print ( pass_through_tensor , values ) : flat_values = [ ] for value in values : if hasattr ( value , '_fields' ) : for field in value . _fields : flat_values . extend ( [ field , _to_str ( getattr ( value , field ) ) ] ) continue if isinstance ( value , ( list , tuple ) ) : for v in value : flat_values . append ( _to_str ( v ) ) continue flat_values . append ( _to_str ( value ) ) return tf . compat . v1 . Print ( pass_through_tensor , flat_values )

Wrapper for tf . Print which supports lists and namedtuples for printing .


def quadrature_scheme_softmaxnormal_gauss_hermite(
    normal_loc, normal_scale, quadrature_size,
    validate_args=False, name=None):
  """Use Gauss-Hermite quadrature to form quadrature on `K - 1` simplex.

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

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

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

  Args:
    normal_loc: `float`-like `Tensor` with shape `[b1, ..., bB, K-1]`, B>=0.
      The location parameter of the Normal used to construct the SoftmaxNormal.
    normal_scale: `float`-like `Tensor`. Broadcastable with `normal_loc`.
      The scale parameter of the Normal used to construct the SoftmaxNormal.
    quadrature_size: Python `int` scalar representing the number of quadrature
      points.
    validate_args: Python `bool`, default `False`. When `True` distribution
      parameters are checked for validity despite possibly degrading runtime
      performance. When `False` invalid inputs may silently render incorrect
      outputs.
    name: Python `str` name prefixed to Ops created by this class.

  Returns:
    grid: Shape `[b1, ..., bB, K, quadrature_size]` `Tensor` representing the
      convex combination of affine parameters for `K` components.
      `grid[..., :, n]` is the `n`-th grid point, living in the `K - 1` simplex.
    probs:  Shape `[b1, ..., bB, K, quadrature_size]` `Tensor` representing the
      associated with each grid point.
  """
  with tf.name_scope(
      name or "quadrature_scheme_softmaxnormal_gauss_hermite"):
    normal_loc = tf.convert_to_tensor(value=normal_loc, name="normal_loc")
    npdt = dtype_util.as_numpy_dtype(normal_loc.dtype)
    normal_scale = tf.convert_to_tensor(
        value=normal_scale, dtype=npdt, name="normal_scale")

    normal_scale = maybe_check_quadrature_param(
        normal_scale, "normal_scale", validate_args)

    grid, probs = np.polynomial.hermite.hermgauss(deg=quadrature_size)
    grid = grid.astype(npdt)
    probs = probs.astype(npdt)
    probs /= np.linalg.norm(probs, ord=1, keepdims=True)
    probs = tf.convert_to_tensor(value=probs, name="probs", dtype=npdt)

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

    return grid, probs

Use Gauss - Hermite quadrature to form quadrature on K - 1 simplex .


def quadrature_scheme_softmaxnormal_quantiles(
    normal_loc, normal_scale, quadrature_size,
    validate_args=False, name=None):
  """Use SoftmaxNormal quantiles to form quadrature on `K - 1` simplex.

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

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

  Args:
    normal_loc: `float`-like `Tensor` with shape `[b1, ..., bB, K-1]`, B>=0.
      The location parameter of the Normal used to construct the SoftmaxNormal.
    normal_scale: `float`-like `Tensor`. Broadcastable with `normal_loc`.
      The scale parameter of the Normal used to construct the SoftmaxNormal.
    quadrature_size: Python `int` scalar representing the number of quadrature
      points.
    validate_args: Python `bool`, default `False`. When `True` distribution
      parameters are checked for validity despite possibly degrading runtime
      performance. When `False` invalid inputs may silently render incorrect
      outputs.
    name: Python `str` name prefixed to Ops created by this class.

  Returns:
    grid: Shape `[b1, ..., bB, K, quadrature_size]` `Tensor` representing the
      convex combination of affine parameters for `K` components.
      `grid[..., :, n]` is the `n`-th grid point, living in the `K - 1` simplex.
    probs:  Shape `[b1, ..., bB, K, quadrature_size]` `Tensor` representing the
      associated with each grid point.
  """
  with tf.name_scope(name or "softmax_normal_grid_and_probs"):
    normal_loc = tf.convert_to_tensor(value=normal_loc, name="normal_loc")
    dt = dtype_util.base_dtype(normal_loc.dtype)
    normal_scale = tf.convert_to_tensor(
        value=normal_scale, dtype=dt, name="normal_scale")

    normal_scale = maybe_check_quadrature_param(
        normal_scale, "normal_scale", validate_args)

    dist = normal.Normal(loc=normal_loc, scale=normal_scale)

    def _get_batch_ndims():
      """Helper to get rank(dist.batch_shape), statically if possible."""
      ndims = tensorshape_util.rank(dist.batch_shape)
      if ndims is None:
        ndims = tf.shape(input=dist.batch_shape_tensor())[0]
      return ndims
    batch_ndims = _get_batch_ndims()

    def _get_final_shape(qs):
      """Helper to build `TensorShape`."""
      bs = tensorshape_util.with_rank_at_least(dist.batch_shape, 1)
      num_components = tf.compat.dimension_value(bs[-1])
      if num_components is not None:
        num_components += 1
      tail = tf.TensorShape([num_components, qs])
      return bs[:-1].concatenate(tail)

    def _compute_quantiles():
      """Helper to build quantiles."""
      # Omit {0, 1} since they might lead to Inf/NaN.
      zero = tf.zeros([], dtype=dist.dtype)
      edges = tf.linspace(zero, 1., quadrature_size + 3)[1:-1]
      # Expand edges so its broadcast across batch dims.
      edges = tf.reshape(
          edges,
          shape=tf.concat(
              [[-1], tf.ones([batch_ndims], dtype=tf.int32)], axis=0))
      quantiles = dist.quantile(edges)
      quantiles = softmax_centered_bijector.SoftmaxCentered().forward(quantiles)
      # Cyclically permute left by one.
      perm = tf.concat([tf.range(1, 1 + batch_ndims), [0]], axis=0)
      quantiles = tf.transpose(a=quantiles, perm=perm)
      tensorshape_util.set_shape(
          quantiles, _get_final_shape(quadrature_size + 1))
      return quantiles
    quantiles = _compute_quantiles()

    # Compute grid as quantile midpoints.
    grid = (quantiles[..., :-1] + quantiles[..., 1:]) / 2.
    # Set shape hints.
    tensorshape_util.set_shape(grid, _get_final_shape(quadrature_size))

    # By construction probs is constant, i.e., `1 / quadrature_size`. This is
    # important, because non-constant probs leads to non-reparameterizable
    # samples.
    probs = tf.fill(
        dims=[quadrature_size], value=1. / tf.cast(quadrature_size, dist.dtype))

    return grid, probs

Use SoftmaxNormal quantiles to form quadrature on K - 1 simplex .


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

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

    if assertions:
      return distribution_util.with_dependencies(assertions, param)
    return param

Helper which checks validity of loc and scale init args .


def determine_batch_event_shapes(grid, endpoint_affine):
  """Helper to infer batch_shape and event_shape."""
  with tf.name_scope("determine_batch_event_shapes"):
    # grid  # shape: [B, k, q]
    # endpoint_affine     # len=k, shape: [B, d, d]
    batch_shape = grid.shape[:-2]
    batch_shape_tensor = tf.shape(input=grid)[:-2]
    event_shape = None
    event_shape_tensor = None

    def _set_event_shape(shape, shape_tensor):
      if event_shape is None:
        return shape, shape_tensor
      return (tf.broadcast_static_shape(event_shape, shape),
              tf.broadcast_dynamic_shape(event_shape_tensor, shape_tensor))

    for aff in endpoint_affine:
      if aff.shift is not None:
        batch_shape = tf.broadcast_static_shape(batch_shape,
                                                aff.shift.shape[:-1])
        batch_shape_tensor = tf.broadcast_dynamic_shape(
            batch_shape_tensor,
            tf.shape(input=aff.shift)[:-1])
        event_shape, event_shape_tensor = _set_event_shape(
            aff.shift.shape[-1:],
            tf.shape(input=aff.shift)[-1:])

      if aff.scale is not None:
        batch_shape = tf.broadcast_static_shape(batch_shape,
                                                aff.scale.batch_shape)
        batch_shape_tensor = tf.broadcast_dynamic_shape(
            batch_shape_tensor, aff.scale.batch_shape_tensor())
        event_shape, event_shape_tensor = _set_event_shape(
            tf.TensorShape([aff.scale.range_dimension]),
            aff.scale.range_dimension_tensor()[tf.newaxis])

    return batch_shape, batch_shape_tensor, event_shape, event_shape_tensor

def determine_batch_event_shapes ( grid , endpoint_affine ) : with tf . name_scope ( "determine_batch_event_shapes" ) : batch_shape = grid . shape [ : - 2 ] batch_shape_tensor = tf . shape ( input = grid ) [ : - 2 ] event_shape = None event_shape_tensor = None def _set_event_shape ( shape , shape_tensor ) : if event_shape is None : return shape , shape_tensor return ( tf . broadcast_static_shape ( event_shape , shape ) , tf . broadcast_dynamic_shape ( event_shape_tensor , shape_tensor ) ) for aff in endpoint_affine : if aff . shift is not None : batch_shape = tf . broadcast_static_shape ( batch_shape , aff . shift . shape [ : - 1 ] ) batch_shape_tensor = tf . broadcast_dynamic_shape ( batch_shape_tensor , tf . shape ( input = aff . shift ) [ : - 1 ] ) event_shape , event_shape_tensor = _set_event_shape ( aff . shift . shape [ - 1 : ] , tf . shape ( input = aff . shift ) [ - 1 : ] ) if aff . scale is not None : batch_shape = tf . broadcast_static_shape ( batch_shape , aff . scale . batch_shape ) batch_shape_tensor = tf . broadcast_dynamic_shape ( batch_shape_tensor , aff . scale . batch_shape_tensor ( ) ) event_shape , event_shape_tensor = _set_event_shape ( tf . TensorShape ( [ aff . scale . range_dimension ] ) , aff . scale . range_dimension_tensor ( ) [ tf . newaxis ] ) return batch_shape , batch_shape_tensor , event_shape , event_shape_tensor

Helper to infer batch_shape and event_shape .


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

Helper which interpolates between two locs .


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

Helper which interpolates between two scales .


def linop_scale(w, op):
  """Creates weighted `LinOp` from existing `LinOp`."""
  # We assume w > 0. (This assumption only relates to the is_* attributes.)
  with tf.name_scope("linop_scale"):
    # TODO(b/35301104): LinearOperatorComposition doesn't combine operators, so
    # special case combinations here. Once it does, this function can be
    # replaced by:
    #     return linop_composition_lib.LinearOperatorComposition([
    #         scaled_identity(w), op])
    def scaled_identity(w):
      return tf.linalg.LinearOperatorScaledIdentity(
          num_rows=op.range_dimension_tensor(),
          multiplier=w,
          is_non_singular=op.is_non_singular,
          is_self_adjoint=op.is_self_adjoint,
          is_positive_definite=op.is_positive_definite)

    if isinstance(op, tf.linalg.LinearOperatorIdentity):
      return scaled_identity(w)
    if isinstance(op, tf.linalg.LinearOperatorScaledIdentity):
      return scaled_identity(w * op.multiplier)
    if isinstance(op, tf.linalg.LinearOperatorDiag):
      return tf.linalg.LinearOperatorDiag(
          diag=w[..., tf.newaxis] * op.diag_part(),
          is_non_singular=op.is_non_singular,
          is_self_adjoint=op.is_self_adjoint,
          is_positive_definite=op.is_positive_definite)
    if isinstance(op, tf.linalg.LinearOperatorLowerTriangular):
      return tf.linalg.LinearOperatorLowerTriangular(
          tril=w[..., tf.newaxis, tf.newaxis] * op.to_dense(),
          is_non_singular=op.is_non_singular,
          is_self_adjoint=op.is_self_adjoint,
          is_positive_definite=op.is_positive_definite)
    raise NotImplementedError(
        "Unsupported Linop type ({})".format(type(op).__name__))

def
linop_scale
(
w
,
op
)
:
with
tf
.
name_scope
(
"linop_scale"
)
:
def
scaled_identity
(
w
)
:
return
tf
.
linalg
.
LinearOperatorScaledIdentity
(
num_rows
=
op
.
range_dimension_tensor
(
)
,
multiplier
=
w
,
is_non_singular
=
op
.
is_non_singular
,
is_self_adjoint
=
op
.
is_self_adjoint
,
is_positive_definite
=
op
.
is_positive_definite
)
if
isinstance
(
op
,
tf
.
linalg
.
LinearOperatorIdentity
)
:
return
scaled_identity
(
w
)
if
isinstance
(
op
,
tf
.
linalg
.
LinearOperatorScaledIdentity
)
:
return
scaled_identity
(
w
*
op
.
multiplier
)
if
isinstance
(
op
,
tf
.
linalg
.
LinearOperatorDiag
)
:
return
tf
.
linalg
.
LinearOperatorDiag
(
diag
=
w
[
...
,
tf
.
newaxis
]
*
op
.
diag_part
(
)
,
is_non_singular
=
op
.
is_non_singular
,
is_self_adjoint
=
op
.
is_self_adjoint
,
is_positive_definite
=
op
.
is_positive_definite
)
if
isinstance
(
op
,
tf
.
linalg
.
LinearOperatorLowerTriangular
)
:
return
tf
.
linalg
.
LinearOperatorLowerTriangular
(
tril
=
w
[
...
,
tf
.
newaxis
,
tf
.
newaxis
]
*
op
.
to_dense
(
)
,
is_non_singular
=
op
.
is_non_singular
,
is_self_adjoint
=
op
.
is_self_adjoint
,
is_positive_definite
=
op
.
is_positive_definite
)
raise
NotImplementedError
(
"Unsupported Linop type ({})"
.
format
(
type
(
op
)
.
__name__
)
)

Creates weighted LinOp from existing LinOp .


def concat_vectors(*args):
  """Concatenates input vectors, statically if possible."""
  args_ = [tf.get_static_value(x) for x in args]
  if any(vec is None for vec in args_):
    return tf.concat(args, axis=0)
  return [val for vec in args_ for val in vec]

def concat_vectors ( * args ) : args_ = [ tf . get_static_value ( x ) for x in args ] if any ( vec is None for vec in args_ ) : return tf . concat ( args , axis = 0 ) return [ val for vec in args_ for val in vec ]

Concatenates input vectors statically if possible .


def _log_vector_matrix(vs, ms):
  """Multiply tensor of vectors by matrices assuming values stored are logs."""

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

def _log_vector_matrix ( vs , ms ) : return tf . reduce_logsumexp ( input_tensor = vs [ ... , tf . newaxis ] + ms , axis = - 2 )

Multiply tensor of vectors by matrices assuming values stored are logs .


def _log_matrix_vector(ms, vs):
  """Multiply tensor of matrices by vectors assuming values stored are logs."""

  return tf.reduce_logsumexp(input_tensor=ms + vs[..., tf.newaxis, :], axis=-1)

def _log_matrix_vector ( ms , vs ) : return tf . reduce_logsumexp ( input_tensor = ms + vs [ ... , tf . newaxis , : ] , axis = - 1 )

Multiply tensor of matrices by vectors assuming values stored are logs .


def _vector_matrix(vs, ms):
  """Multiply tensor of vectors by matrices."""

  return tf.reduce_sum(input_tensor=vs[..., tf.newaxis] * ms, axis=-2)

def _vector_matrix ( vs , ms ) : return tf . reduce_sum ( input_tensor = vs [ ... , tf . newaxis ] * ms , axis = - 2 )

Multiply tensor of vectors by matrices .


def _extract_log_probs(num_states, dist):
  """Tabulate log probabilities from a batch of distributions."""

  states = tf.reshape(tf.range(num_states),
                      tf.concat([[num_states],
                                 tf.ones_like(dist.batch_shape_tensor())],
                                axis=0))
  return distribution_util.move_dimension(dist.log_prob(states), 0, -1)

def _extract_log_probs ( num_states , dist ) : states = tf . reshape ( tf . range ( num_states ) , tf . concat ( [ [ num_states ] , tf . ones_like ( dist . batch_shape_tensor ( ) ) ] , axis = 0 ) ) return distribution_util . move_dimension ( dist . log_prob ( states ) , 0 , - 1 )

Tabulate log probabilities from a batch of distributions .


def _marginal_hidden_probs(self):
    """Compute marginal pdf for each individual observable."""

    initial_log_probs = tf.broadcast_to(self._log_init,
                                        tf.concat([self.batch_shape_tensor(),
                                                   [self._num_states]],
                                                  axis=0))
    # initial_log_probs :: batch_shape num_states

    if self._num_steps > 1:
      transition_log_probs = self._log_trans

      def forward_step(log_probs, _):
        return _log_vector_matrix(log_probs, transition_log_probs)

      dummy_index = tf.zeros(self._num_steps - 1, dtype=tf.float32)

      forward_log_probs = tf.scan(forward_step, dummy_index,
                                  initializer=initial_log_probs,
                                  name="forward_log_probs")

      forward_log_probs = tf.concat([[initial_log_probs], forward_log_probs],
                                    axis=0)
    else:
      forward_log_probs = initial_log_probs[tf.newaxis, ...]

    # returns :: num_steps batch_shape num_states

    return tf.exp(forward_log_probs)

def _marginal_hidden_probs ( self ) : initial_log_probs = tf . broadcast_to ( self . _log_init , tf . concat ( [ self . batch_shape_tensor ( ) , [ self . _num_states ] ] , axis = 0 ) ) if self . _num_steps > 1 : transition_log_probs = self . _log_trans def forward_step ( log_probs , _ ) : return _log_vector_matrix ( log_probs , transition_log_probs ) dummy_index = tf . zeros ( self . _num_steps - 1 , dtype = tf . float32 ) forward_log_probs = tf . scan ( forward_step , dummy_index , initializer = initial_log_probs , name = "forward_log_probs" ) forward_log_probs = tf . concat ( [ [ initial_log_probs ] , forward_log_probs ] , axis = 0 ) else : forward_log_probs = initial_log_probs [ tf . newaxis , ... ] return tf . exp ( forward_log_probs )

Compute marginal pdf for each individual observable .


def posterior_marginals(self, observations, name=None):
    """Compute marginal posterior distribution for each state.

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

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

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

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

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

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

    with tf.name_scope(name or "posterior_marginals"):
      with tf.control_dependencies(self._runtime_assertions):
        observation_tensor_shape = tf.shape(input=observations)

        with self._observation_shape_preconditions(observation_tensor_shape):
          observation_batch_shape = observation_tensor_shape[
              :-1 - self._underlying_event_rank]
          observation_event_shape = observation_tensor_shape[
              -1 - self._underlying_event_rank:]

          batch_shape = tf.broadcast_dynamic_shape(observation_batch_shape,
                                                   self.batch_shape_tensor())
          log_init = tf.broadcast_to(self._log_init,
                                     tf.concat([batch_shape,
                                                [self._num_states]],
                                               axis=0))
          log_transition = self._log_trans

          observations = tf.broadcast_to(observations,
                                         tf.concat([batch_shape,
                                                    observation_event_shape],
                                                   axis=0))
          observation_rank = tf.rank(observations)
          underlying_event_rank = self._underlying_event_rank
          observations = distribution_util.move_dimension(
              observations, observation_rank - underlying_event_rank - 1, 0)
          observations = tf.expand_dims(
              observations,
              observation_rank - underlying_event_rank)
          observation_log_probs = self._observation_distribution.log_prob(
              observations)

          log_adjoint_prob = tf.zeros_like(log_init)

          def forward_step(log_previous_step, log_prob_observation):
            return _log_vector_matrix(log_previous_step,
                                      log_transition) + log_prob_observation

          log_prob = log_init + observation_log_probs[0]

          forward_log_probs = tf.scan(forward_step, observation_log_probs[1:],
                                      initializer=log_prob,
                                      name="forward_log_probs")

          forward_log_probs = tf.concat([[log_prob], forward_log_probs], axis=0)

          def backward_step(log_previous_step, log_prob_observation):
            return _log_matrix_vector(log_transition,
                                      log_prob_observation + log_previous_step)

          backward_log_adjoint_probs = tf.scan(
              backward_step,
              observation_log_probs[1:],
              initializer=log_adjoint_prob,
              reverse=True,
              name="backward_log_adjoint_probs")

          total_log_prob = tf.reduce_logsumexp(
              input_tensor=forward_log_probs[-1], axis=-1)

          backward_log_adjoint_probs = tf.concat([backward_log_adjoint_probs,
                                                  [log_adjoint_prob]], axis=0)

          log_likelihoods = forward_log_probs + backward_log_adjoint_probs

          marginal_log_probs = distribution_util.move_dimension(
              log_likelihoods - total_log_prob[..., tf.newaxis], 0, -2)

          return categorical.Categorical(logits=marginal_log_probs)

def posterior_marginals ( self , observations , name = None ) : with tf . name_scope ( name or "posterior_marginals" ) : with tf . control_dependencies ( self . _runtime_assertions ) : observation_tensor_shape = tf . shape ( input = observations ) with self . _observation_shape_preconditions ( observation_tensor_shape ) : observation_batch_shape = observation_tensor_shape [ : - 1 - self . _underlying_event_rank ] observation_event_shape = observation_tensor_shape [ - 1 - self . _underlying_event_rank : ] batch_shape = tf . broadcast_dynamic_shape ( observation_batch_shape , self . batch_shape_tensor ( ) ) log_init = tf . broadcast_to ( self . _log_init , tf . concat ( [ batch_shape , [ self . _num_states ] ] , axis = 0 ) ) log_transition = self . _log_trans observations = tf . broadcast_to ( observations , tf . concat ( [ batch_shape , observation_event_shape ] , axis = 0 ) ) observation_rank = tf . rank ( observations ) underlying_event_rank = self . _underlying_event_rank observations = distribution_util . move_dimension ( observations , observation_rank - underlying_event_rank - 1 , 0 ) observations = tf . expand_dims ( observations , observation_rank - underlying_event_rank ) observation_log_probs = self . _observation_distribution . log_prob ( observations ) log_adjoint_prob = tf . zeros_like ( log_init ) def forward_step ( log_previous_step , log_prob_observation ) : return _log_vector_matrix ( log_previous_step , log_transition ) + log_prob_observation log_prob = log_init + observation_log_probs [ 0 ] forward_log_probs = tf . scan ( forward_step , observation_log_probs [ 1 : ] , initializer = log_prob , name = "forward_log_probs" ) forward_log_probs = tf . concat ( [ [ log_prob ] , forward_log_probs ] , axis = 0 ) def backward_step ( log_previous_step , log_prob_observation ) : return _log_matrix_vector ( log_transition , log_prob_observation + log_previous_step ) backward_log_adjoint_probs = tf . scan ( backward_step , observation_log_probs [ 1 : ] , initializer = log_adjoint_prob , reverse = True , name = "backward_log_adjoint_probs" ) total_log_prob = tf . reduce_logsumexp ( input_tensor = forward_log_probs [ - 1 ] , axis = - 1 ) backward_log_adjoint_probs = tf . concat ( [ backward_log_adjoint_probs , [ log_adjoint_prob ] ] , axis = 0 ) log_likelihoods = forward_log_probs + backward_log_adjoint_probs marginal_log_probs = distribution_util . move_dimension ( log_likelihoods - total_log_prob [ ... , tf . newaxis ] , 0 , - 2 ) return categorical . Categorical ( logits = marginal_log_probs )

Compute marginal posterior distribution for each state .


def posterior_mode(self, observations, name=None):
    """Compute maximum likelihood sequence of hidden states.

    When this function is provided with a sequence of observations
    `x[0], ..., x[num_steps - 1]`, it returns the sequence of hidden
    states `z[0], ..., z[num_steps - 1]`, drawn from the underlying
    Markov chain, that is most likely to yield those observations.

    It uses the [Viterbi algorithm](
    https://en.wikipedia.org/wiki/Viterbi_algorithm).

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

    Note: if there isn't a unique most likely sequence then one
    of the equally most likely sequences is chosen.

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

    Returns:
      posterior_mode: A `Tensor` representing the most likely sequence of hidden
        states. The rightmost dimension of this tensor will equal the
        `num_steps` parameter providing one hidden state for each step. The
        other dimensions are those of the batch.

    Raises:
      ValueError: if the `observations` tensor does not consist of
      sequences of `num_steps` observations.

    #### Examples

    ```python
    tfd = tfp.distributions

    # A simple weather model.

    # Represent a cold day with 0 and a hot day with 1.
    # Suppose the first day of a sequence has a 0.8 chance of being cold.

    initial_distribution = tfd.Categorical(probs=[0.8, 0.2])

    # Suppose a cold day has a 30% chance of being followed by a hot day
    # and a hot day has a 20% chance of being followed by a cold day.

    transition_distribution = tfd.Categorical(probs=[[0.7, 0.3],
                                                     [0.2, 0.8]])

    # Suppose additionally that on each day the temperature is
    # normally distributed with mean and standard deviation 0 and 5 on
    # a cold day and mean and standard deviation 15 and 10 on a hot day.

    observation_distribution = tfd.Normal(loc=[0., 15.], scale=[5., 10.])

    # This gives the hidden Markov model:

    model = tfd.HiddenMarkovModel(
        initial_distribution=initial_distribution,
        transition_distribution=transition_distribution,
        observation_distribution=observation_distribution,
        num_steps=7)

    # Suppose we observe gradually rising temperatures over a week:
    temps = [-2., 0., 2., 4., 6., 8., 10.]

    # We can now compute the most probable sequence of hidden states:

    model.posterior_mode(temps)

    # The result is [0 0 0 0 0 1 1] telling us that the transition
    # from "cold" to "hot" most likely happened between the
    # 5th and 6th days.
    ```
    """

    with tf.name_scope(name or "posterior_mode"):
      with tf.control_dependencies(self._runtime_assertions):
        observation_tensor_shape = tf.shape(input=observations)

        with self._observation_shape_preconditions(observation_tensor_shape):
          observation_batch_shape = observation_tensor_shape[
              :-1 - self._underlying_event_rank]
          observation_event_shape = observation_tensor_shape[
              -1 - self._underlying_event_rank:]

          batch_shape = tf.broadcast_dynamic_shape(observation_batch_shape,
                                                   self.batch_shape_tensor())
          log_init = tf.broadcast_to(self._log_init,
                                     tf.concat([batch_shape,
                                                [self._num_states]],
                                               axis=0))

          observations = tf.broadcast_to(observations,
                                         tf.concat([batch_shape,
                                                    observation_event_shape],
                                                   axis=0))
          observation_rank = tf.rank(observations)
          underlying_event_rank = self._underlying_event_rank
          observations = distribution_util.move_dimension(
              observations, observation_rank - underlying_event_rank - 1, 0)

          # We need to compute the probability of each observation for
          # each possible state.
          # This requires inserting an extra index just before the
          # observation event indices that will be broadcast with the
          # last batch index in `observation_distribution`.
          observations = tf.expand_dims(
              observations,
              observation_rank - underlying_event_rank)
          observation_log_probs = self._observation_distribution.log_prob(
              observations)

          log_prob = log_init + observation_log_probs[0]

          if self._num_steps == 1:
            most_likely_end = tf.argmax(input=log_prob, axis=-1)
            return most_likely_end[..., tf.newaxis]

          def forward_step(previous_step_pair, log_prob_observation):
            log_prob_previous = previous_step_pair[0]
            log_prob = (log_prob_previous[..., tf.newaxis] +
                        self._log_trans +
                        log_prob_observation[..., tf.newaxis, :])
            most_likely_given_successor = tf.argmax(input=log_prob, axis=-2)
            max_log_p_given_successor = tf.reduce_max(input_tensor=log_prob,
                                                      axis=-2)
            return (max_log_p_given_successor, most_likely_given_successor)

          forward_log_probs, all_most_likely_given_successor = tf.scan(
              forward_step,
              observation_log_probs[1:],
              initializer=(log_prob,
                           tf.zeros(tf.shape(input=log_init), dtype=tf.int64)),
              name="forward_log_probs")

          most_likely_end = tf.argmax(input=forward_log_probs[-1], axis=-1)

          # We require the operation that gives C from A and B where
          # C[i...j] = A[i...j, B[i...j]]
          # and A = most_likely_given_successor
          #     B = most_likely_successor.
          # tf.gather requires indices of known shape so instead we use
          # reduction with tf.one_hot(B) to pick out elements from B
          def backward_step(most_likely_successor, most_likely_given_successor):
            return tf.reduce_sum(
                input_tensor=(most_likely_given_successor *
                              tf.one_hot(most_likely_successor,
                                         self._num_states,
                                         dtype=tf.int64)),
                axis=-1)

          backward_scan = tf.scan(
              backward_step,
              all_most_likely_given_successor,
              most_likely_end,
              reverse=True)
          most_likely_sequences = tf.concat([backward_scan, [most_likely_end]],
                                            axis=0)
          return distribution_util.move_dimension(most_likely_sequences, 0, -1)

def posterior_mode ( self , observations , name = None ) : with tf . name_scope ( name or "posterior_mode" ) : with tf . control_dependencies ( self . _runtime_assertions ) : observation_tensor_shape = tf . shape ( input = observations ) with self . _observation_shape_preconditions ( observation_tensor_shape ) : observation_batch_shape = observation_tensor_shape [ : - 1 - self . _underlying_event_rank ] observation_event_shape = observation_tensor_shape [ - 1 - self . _underlying_event_rank : ] batch_shape = tf . broadcast_dynamic_shape ( observation_batch_shape , self . batch_shape_tensor ( ) ) log_init = tf . broadcast_to ( self . _log_init , tf . concat ( [ batch_shape , [ self . _num_states ] ] , axis = 0 ) ) observations = tf . broadcast_to ( observations , tf . concat ( [ batch_shape , observation_event_shape ] , axis = 0 ) ) observation_rank = tf . rank ( observations ) underlying_event_rank = self . _underlying_event_rank observations = distribution_util . move_dimension ( observations , observation_rank - underlying_event_rank - 1 , 0 ) observations = tf . expand_dims ( observations , observation_rank - underlying_event_rank ) observation_log_probs = self . _observation_distribution . log_prob ( observations ) log_prob = log_init + observation_log_probs [ 0 ] if self . _num_steps == 1 : most_likely_end = tf . argmax ( input = log_prob , axis = - 1 ) return most_likely_end [ ... , tf . newaxis ] def forward_step ( previous_step_pair , log_prob_observation ) : log_prob_previous = previous_step_pair [ 0 ] log_prob = ( log_prob_previous [ ... , tf . newaxis ] + self . _log_trans + log_prob_observation [ ... , tf . newaxis , : ] ) most_likely_given_successor = tf . argmax ( input = log_prob , axis = - 2 ) max_log_p_given_successor = tf . reduce_max ( input_tensor = log_prob , axis = - 2 ) return ( max_log_p_given_successor , most_likely_given_successor ) forward_log_probs , all_most_likely_given_successor = tf . scan ( forward_step , observation_log_probs [ 1 : ] , initializer = ( log_prob , tf . zeros ( tf . shape ( input = log_init ) , dtype = tf . int64 ) ) , name = "forward_log_probs" ) most_likely_end = tf . argmax ( input = forward_log_probs [ - 1 ] , axis = - 1 ) def backward_step ( most_likely_successor , most_likely_given_successor ) : return tf . reduce_sum ( input_tensor = ( most_likely_given_successor * tf . one_hot ( most_likely_successor , self . _num_states , dtype = tf . int64 ) ) , axis = - 1 ) backward_scan = tf . scan ( backward_step , all_most_likely_given_successor , most_likely_end , reverse = True ) most_likely_sequences = tf . concat ( [ backward_scan , [ most_likely_end ] ] , axis = 0 ) return distribution_util . move_dimension ( most_likely_sequences , 0 , - 1 )

Compute maximum likelihood sequence of hidden states .


def _choose_random_direction(current_state_parts, batch_rank, seed=None):
  """Chooses a random direction in the event space."""
  seed_gen = distributions.SeedStream(seed, salt='_choose_random_direction')
  # Chooses the random directions across each of the input components.
  rnd_direction_parts = [
      tf.random.normal(
          tf.shape(input=current_state_part), dtype=tf.float32, seed=seed_gen())
      for current_state_part in current_state_parts
  ]

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

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

  return rnd_direction_parts

def _choose_random_direction ( current_state_parts , batch_rank , seed = None ) : seed_gen = distributions . SeedStream ( seed , salt = '_choose_random_direction' ) rnd_direction_parts = [ tf . random . normal ( tf . shape ( input = current_state_part ) , dtype = tf . float32 , seed = seed_gen ( ) ) for current_state_part in current_state_parts ] sum_squares = sum ( tf . reduce_sum ( input_tensor = rnd_direction ** 2. , axis = tf . range ( batch_rank , tf . rank ( rnd_direction ) ) , keepdims = True ) for rnd_direction in rnd_direction_parts ) rnd_direction_parts = [ rnd_direction / tf . sqrt ( sum_squares ) for rnd_direction in rnd_direction_parts ] return rnd_direction_parts

Chooses a random direction in the event space .


def _sample_next(target_log_prob_fn,
                 current_state_parts,
                 step_sizes,
                 max_doublings,
                 current_target_log_prob,
                 batch_rank,
                 seed=None,
                 name=None):
  """Applies a single iteration of slice sampling update.

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

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

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

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

    components = [
        tf.reduce_sum(
            input_tensor=(dirn_part / step_size)**2, axis=reduce_axes[i])
        for i, (step_size, dirn_part) in enumerate(zip(step_sizes, direction))
    ]
    step_size = tf.math.rsqrt(tf.add_n(components))
    # Computes the rank of a tensor. Uses the static rank if possible.
    def _get_rank(x):
      return (len(x.shape.as_list()) if x.shape.dims is not None
              else tf.rank(x))
    state_part_ranks = [_get_rank(part) for part in current_state_parts]

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

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

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

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

      state_parts = [state_part + padded_alpha
                     * direction_part for state_part, direction_part,
                     padded_alpha in
                     zip(current_state_parts, direction, padded_alphas)]
      return state_parts

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

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

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

    alpha_init = tf.zeros_like(current_target_log_prob,
                               dtype=current_state_parts[0].dtype.base_dtype)
    [
        next_alpha,
        next_target_log_prob,
        bounds_satisfied,
        upper_bounds,
        lower_bounds
    ] = ssu.slice_sampler_one_dim(projected_target_log_prob_fn,
                                  x_initial=alpha_init,
                                  max_doublings=max_doublings,
                                  step_size=step_size, seed=seed)
    return [
        _step_along_direction(next_alpha),
        next_target_log_prob,
        bounds_satisfied,
        direction,
        upper_bounds,
        lower_bounds
    ]

def _sample_next ( target_log_prob_fn , current_state_parts , step_sizes , max_doublings , current_target_log_prob , batch_rank , seed = None , name = None ) : with tf . compat . v1 . name_scope ( name , 'sample_next' , [ current_state_parts , step_sizes , max_doublings , current_target_log_prob , batch_rank ] ) : direction = _choose_random_direction ( current_state_parts , batch_rank = batch_rank , seed = seed ) reduce_axes = [ tf . range ( batch_rank , tf . rank ( dirn_part ) ) for dirn_part in direction ] components = [ tf . reduce_sum ( input_tensor = ( dirn_part / step_size ) ** 2 , axis = reduce_axes [ i ] ) for i , ( step_size , dirn_part ) in enumerate ( zip ( step_sizes , direction ) ) ] step_size = tf . math . rsqrt ( tf . add_n ( components ) ) def _get_rank ( x ) : return ( len ( x . shape . as_list ( ) ) if x . shape . dims is not None else tf . rank ( x ) ) state_part_ranks = [ _get_rank ( part ) for part in current_state_parts ] def _step_along_direction ( alpha ) : padded_alphas = [ _right_pad ( alpha , final_rank = part_rank ) for part_rank in state_part_ranks ] state_parts = [ state_part + padded_alpha * direction_part for state_part , direction_part , padded_alpha in zip ( current_state_parts , direction , padded_alphas ) ] return state_parts def projected_target_log_prob_fn ( alpha ) : return target_log_prob_fn ( * _step_along_direction ( alpha ) ) alpha_init = tf . zeros_like ( current_target_log_prob , dtype = current_state_parts [ 0 ] . dtype . base_dtype ) [ next_alpha , next_target_log_prob , bounds_satisfied , upper_bounds , lower_bounds ] = ssu . slice_sampler_one_dim ( projected_target_log_prob_fn , x_initial = alpha_init , max_doublings = max_doublings , step_size = step_size , seed = seed ) return [ _step_along_direction ( next_alpha ) , next_target_log_prob , bounds_satisfied , direction , upper_bounds , lower_bounds ]

Applies a single iteration of slice sampling update .


def _maybe_call_fn(fn,
                   fn_arg_list,
                   fn_result=None,
                   description='target_log_prob'):
  """Helper which computes `fn_result` if needed."""
  fn_arg_list = (list(fn_arg_list) if mcmc_util.is_list_like(fn_arg_list)
                 else [fn_arg_list])
  if fn_result is None:
    fn_result = fn(*fn_arg_list)
  if not fn_result.dtype.is_floating:
    raise TypeError('`{}` must be a `Tensor` with `float` `dtype`.'.format(
        description))
  return fn_result

Helper which computes fn_result if needed .


def _right_pad(x, final_rank):
  """Pads the shape of x to the right to be of rank final_rank.

  Expands the dims of `x` to the right such that its rank is equal to
  final_rank. For example, if `x` is of shape [1, 5, 7, 2] and `final_rank` is
  7, we return padded_x, which is of shape [1, 5, 7, 2, 1, 1, 1].

  Args:
    x: The tensor whose shape is to be padded.
    final_rank: Scalar int32 `Tensor` or Python `int`. The desired rank of x.

  Returns:
    padded_x: A tensor of rank final_rank.
  """
  padded_shape = tf.concat(
      [tf.shape(input=x),
       tf.ones(final_rank - tf.rank(x), dtype=tf.int32)],
      axis=0)
  static_padded_shape = None
  if x.shape.is_fully_defined() and isinstance(final_rank, int):
    static_padded_shape = x.shape.as_list()
    extra_dims = final_rank - len(static_padded_shape)
    static_padded_shape.extend([1] * extra_dims)

  padded_x = tf.reshape(x, static_padded_shape or padded_shape)
  return padded_x

def _right_pad ( x , final_rank ) : padded_shape = tf . concat ( [ tf . shape ( input = x ) , tf . ones ( final_rank - tf . rank ( x ) , dtype = tf . int32 ) ] , axis = 0 ) static_padded_shape = None if x . shape . is_fully_defined ( ) and isinstance ( final_rank , int ) : static_padded_shape = x . shape . as_list ( ) extra_dims = final_rank - len ( static_padded_shape ) static_padded_shape . extend ( [ 1 ] * extra_dims ) padded_x = tf . reshape ( x , static_padded_shape or padded_shape ) return padded_x

Pads the shape of x to the right to be of rank final_rank .


def one_step(self, current_state, previous_kernel_results):
    """Runs one iteration of Slice Sampler.

    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`.
      TypeError: if `not target_log_prob.dtype.is_floating`.
    """
    with tf.compat.v1.name_scope(
        name=mcmc_util.make_name(self.name, 'slice', 'one_step'),
        values=[
            self.step_size, self.max_doublings, self._seed_stream,
            current_state, previous_kernel_results.target_log_prob
        ]):
      with tf.compat.v1.name_scope('initialize'):
        [
            current_state_parts,
            step_sizes,
            current_target_log_prob
        ] = _prepare_args(
            self.target_log_prob_fn,
            current_state,
            self.step_size,
            previous_kernel_results.target_log_prob,
            maybe_expand=True)

        max_doublings = tf.convert_to_tensor(
            value=self.max_doublings, dtype=tf.int32, name='max_doublings')

      independent_chain_ndims = distribution_util.prefer_static_rank(
          current_target_log_prob)

      [
          next_state_parts,
          next_target_log_prob,
          bounds_satisfied,
          direction,
          upper_bounds,
          lower_bounds
      ] = _sample_next(
          self.target_log_prob_fn,
          current_state_parts,
          step_sizes,
          max_doublings,
          current_target_log_prob,
          independent_chain_ndims,
          seed=self._seed_stream()
      )

      def maybe_flatten(x):
        return x if mcmc_util.is_list_like(current_state) else x[0]

      return [
          maybe_flatten(next_state_parts),
          SliceSamplerKernelResults(
              target_log_prob=next_target_log_prob,
              bounds_satisfied=bounds_satisfied,
              direction=direction,
              upper_bounds=upper_bounds,
              lower_bounds=lower_bounds
          ),
      ]

def one_step ( self , current_state , previous_kernel_results ) : with tf . compat . v1 . name_scope ( name = mcmc_util . make_name ( self . name , 'slice' , 'one_step' ) , values = [ self . step_size , self . max_doublings , self . _seed_stream , current_state , previous_kernel_results . target_log_prob ] ) : with tf . compat . v1 . name_scope ( 'initialize' ) : [ current_state_parts , step_sizes , current_target_log_prob ] = _prepare_args ( self . target_log_prob_fn , current_state , self . step_size , previous_kernel_results . target_log_prob , maybe_expand = True ) max_doublings = tf . convert_to_tensor ( value = self . max_doublings , dtype = tf . int32 , name = 'max_doublings' ) independent_chain_ndims = distribution_util . prefer_static_rank ( current_target_log_prob ) [ next_state_parts , next_target_log_prob , bounds_satisfied , direction , upper_bounds , lower_bounds ] = _sample_next ( self . target_log_prob_fn , current_state_parts , step_sizes , max_doublings , current_target_log_prob , independent_chain_ndims , seed = self . _seed_stream ( ) ) def maybe_flatten ( x ) : return x if mcmc_util . is_list_like ( current_state ) else x [ 0 ] return [ maybe_flatten ( next_state_parts ) , SliceSamplerKernelResults ( target_log_prob = next_target_log_prob , bounds_satisfied = bounds_satisfied , direction = direction , upper_bounds = upper_bounds , lower_bounds = lower_bounds ) , ]


def _build_trainable_posterior(param, initial_loc_fn):
  """Built a transformed-normal variational dist over a parameter's support."""
  loc = tf.compat.v1.get_variable(
      param.name + '_loc',
      initializer=lambda: initial_loc_fn(param),
      dtype=param.prior.dtype,
      use_resource=True)
  scale = tf.nn.softplus(
      tf.compat.v1.get_variable(
          param.name + '_scale',
          initializer=lambda: -4 * tf.ones_like(initial_loc_fn(param)),
          dtype=param.prior.dtype,
          use_resource=True))

  q = tfd.Normal(loc=loc, scale=scale)

  # Ensure the `event_shape` of the variational distribution matches the
  # parameter.
  if (param.prior.event_shape.ndims is None
      or param.prior.event_shape.ndims > 0):
    q = tfd.Independent(
        q, reinterpreted_batch_ndims=param.prior.event_shape.ndims)

  # Transform to constrained parameter space.
  return tfd.TransformedDistribution(q, param.bijector)

def _build_trainable_posterior ( param , initial_loc_fn ) : loc = tf . compat . v1 . get_variable ( param . name + '_loc' , initializer = lambda : initial_loc_fn ( param ) , dtype = param . prior . dtype , use_resource = True ) scale = tf . nn . softplus ( tf . compat . v1 . get_variable ( param . name + '_scale' , initializer = lambda : - 4 * tf . ones_like ( initial_loc_fn ( param ) ) , dtype = param . prior . dtype , use_resource = True ) ) q = tfd . Normal ( loc = loc , scale = scale ) if ( param . prior . event_shape . ndims is None or param . prior . event_shape . ndims > 0 ) : q = tfd . Independent ( q , reinterpreted_batch_ndims = param . prior . event_shape . ndims ) return tfd . TransformedDistribution ( q , param . bijector )

Built a transformed - normal variational dist over a parameter s support .