def random_walk_normal_fn(scale=1., name=None):
  """Returns a callable that adds a random normal perturbation to the input.

  This function returns a callable that accepts a Python `list` of `Tensor`s of
  any shapes and `dtypes`  representing the state parts of the `current_state`
  and a random seed. The supplied argument `scale` must be a `Tensor` or Python
  `list` of `Tensor`s representing the scale of the generated
  proposal. `scale` must broadcast with the state parts of `current_state`.
  The callable adds a sample from a zero-mean normal distribution with the
  supplied scales to each state part and returns a same-type `list` of `Tensor`s
  as the state parts of `current_state`.

  Args:
    scale: a `Tensor` or Python `list` of `Tensor`s of any shapes and `dtypes`
      controlling the scale of the normal proposal distribution.
    name: Python `str` name prefixed to Ops created by this function.
        Default value: 'random_walk_normal_fn'.

  Returns:
    random_walk_normal_fn: A callable accepting a Python `list` of `Tensor`s
      representing the state parts of the `current_state` and an `int`
      representing the random seed to be used to generate the proposal. The
      callable returns the same-type `list` of `Tensor`s as the input and
      represents the proposal for the RWM algorithm.
  """
  def _fn(state_parts, seed):
    """Adds a normal perturbation to the input state.

    Args:
      state_parts: A list of `Tensor`s of any shape and real dtype representing
        the state parts of the `current_state` of the Markov chain.
      seed: `int` or None. The random seed for this `Op`. If `None`, no seed is
        applied.
        Default value: `None`.

    Returns:
      perturbed_state_parts: A Python `list` of The `Tensor`s. Has the same
        shape and type as the `state_parts`.

    Raises:
      ValueError: if `scale` does not broadcast with `state_parts`.
    """
    with tf.compat.v1.name_scope(
        name, 'random_walk_normal_fn', values=[state_parts, scale, seed]):
      scales = scale if mcmc_util.is_list_like(scale) else [scale]
      if len(scales) == 1:
        scales *= len(state_parts)
      if len(state_parts) != len(scales):
        raise ValueError('`scale` must broadcast with `state_parts`.')
      seed_stream = distributions.SeedStream(seed, salt='RandomWalkNormalFn')
      next_state_parts = [
          tf.random.normal(
              mean=state_part,
              stddev=scale_part,
              shape=tf.shape(input=state_part),
              dtype=state_part.dtype.base_dtype,
              seed=seed_stream())
          for scale_part, state_part in zip(scales, state_parts)
      ]

      return next_state_parts
  return _fn

Returns a callable that adds a random normal perturbation to the input .


def random_walk_uniform_fn(scale=1., name=None):
  """Returns a callable that adds a random uniform perturbation to the input.

  For more details on `random_walk_uniform_fn`, see
  `random_walk_normal_fn`. `scale` might
  be a `Tensor` or a list of `Tensor`s that should broadcast with state parts
  of the `current_state`. The generated uniform perturbation is sampled as a
  uniform point on the rectangle `[-scale, scale]`.

  Args:
    scale: a `Tensor` or Python `list` of `Tensor`s of any shapes and `dtypes`
      controlling the upper and lower bound of the uniform proposal
      distribution.
    name: Python `str` name prefixed to Ops created by this function.
        Default value: 'random_walk_uniform_fn'.

  Returns:
    random_walk_uniform_fn: A callable accepting a Python `list` of `Tensor`s
      representing the state parts of the `current_state` and an `int`
      representing the random seed used to generate the proposal. The callable
      returns the same-type `list` of `Tensor`s as the input and represents the
      proposal for the RWM algorithm.
  """
  def _fn(state_parts, seed):
    """Adds a uniform perturbation to the input state.

    Args:
      state_parts: A list of `Tensor`s of any shape and real dtype representing
        the state parts of the `current_state` of the Markov chain.
      seed: `int` or None. The random seed for this `Op`. If `None`, no seed is
        applied.
        Default value: `None`.

    Returns:
      perturbed_state_parts: A Python `list` of The `Tensor`s. Has the same
        shape and type as the `state_parts`.

    Raises:
      ValueError: if `scale` does not broadcast with `state_parts`.
    """
    with tf.compat.v1.name_scope(
        name, 'random_walk_uniform_fn', values=[state_parts, scale, seed]):
      scales = scale if mcmc_util.is_list_like(scale) else [scale]
      if len(scales) == 1:
        scales *= len(state_parts)
      if len(state_parts) != len(scales):
        raise ValueError('`scale` must broadcast with `state_parts`.')
      seed_stream = distributions.SeedStream(seed, salt='RandomWalkUniformFn')
      next_state_parts = [
          tf.random.uniform(
              minval=state_part - scale_part,
              maxval=state_part + scale_part,
              shape=tf.shape(input=state_part),
              dtype=state_part.dtype.base_dtype,
              seed=seed_stream())
          for scale_part, state_part in zip(scales, state_parts)
      ]
      return next_state_parts
  return _fn

Returns a callable that adds a random uniform perturbation to the input .


def _expand_to_event_rank(self, x):
    """Expand the rank of x up to static_event_rank times for broadcasting.

    The static event rank was checked to not be None at construction time.

    Args:
      x: A tensor to expand.
    Returns:
      The expanded tensor.
    """
    expanded_x = x
    for _ in range(tensorshape_util.rank(self.event_shape)):
      expanded_x = tf.expand_dims(expanded_x, -1)
    return expanded_x

def _expand_to_event_rank ( self , x ) : expanded_x = x for _ in range ( tensorshape_util . rank ( self . event_shape ) ) : expanded_x = tf . expand_dims ( expanded_x , - 1 ) return expanded_x

Expand the rank of x up to static_event_rank times for broadcasting .


def entropy_lower_bound(self, name="entropy_lower_bound"):
    r"""A lower bound on the entropy of this mixture model.

    The bound below is not always very tight, and its usefulness depends
    on the mixture probabilities and the components in use.

    A lower bound is useful for ELBO when the `Mixture` is the variational
    distribution:

    \\(
    \log p(x) >= ELBO = \int q(z) \log p(x, z) dz + H[q]
    \\)

    where \\( p \\) is the prior distribution, \\( q \\) is the variational,
    and \\( H[q] \\) is the entropy of \\( q \\). If there is a lower bound
    \\( G[q] \\) such that \\( H[q] \geq G[q] \\) then it can be used in
    place of \\( H[q] \\).

    For a mixture of distributions \\( q(Z) = \sum_i c_i q_i(Z) \\) with
    \\( \sum_i c_i = 1 \\), by the concavity of \\( f(x) = -x \log x \\), a
    simple lower bound is:

    \\(
    \begin{align}
    H[q] & = - \int q(z) \log q(z) dz \\\
       & = - \int (\sum_i c_i q_i(z)) \log(\sum_i c_i q_i(z)) dz \\\
       & \geq - \sum_i c_i \int q_i(z) \log q_i(z) dz \\\
       & = \sum_i c_i H[q_i]
    \end{align}
    \\)

    This is the term we calculate below for \\( G[q] \\).

    Args:
      name: A name for this operation (optional).

    Returns:
      A lower bound on the Mixture's entropy.
    """
    with self._name_scope(name):
      with tf.control_dependencies(self._assertions):
        distribution_entropies = [d.entropy() for d in self.components]
        cat_probs = self._cat_probs(log_probs=False)
        partial_entropies = [
            c_p * m for (c_p, m) in zip(cat_probs, distribution_entropies)
        ]
        # These are all the same shape by virtue of matching batch_shape
        return tf.add_n(partial_entropies)

r A lower bound on the entropy of this mixture model .


def _cat_probs(self, log_probs):
    """Get a list of num_components batchwise probabilities."""
    which_softmax = tf.nn.log_softmax if log_probs else tf.nn.softmax
    cat_probs = which_softmax(self.cat.logits)
    cat_probs = tf.unstack(cat_probs, num=self.num_components, axis=-1)
    return cat_probs

def _cat_probs ( self , log_probs ) : which_softmax = tf . nn . log_softmax if log_probs else tf . nn . softmax cat_probs = which_softmax ( self . cat . logits ) cat_probs = tf . unstack ( cat_probs , num = self . num_components , axis = - 1 ) return cat_probs

Get a list of num_components batchwise probabilities .


def _maybe_validate_args(outcomes, logits, probs, validate_args):
  """Validate `outcomes`, `logits` and `probs`'s shapes."""
  assertions = []

  def validate_equal_last_dim(tensor_a, tensor_b, message):
    if tensor_a.shape.is_fully_defined() and tensor_b.shape.is_fully_defined():
      if tensor_a.shape[-1] != tensor_b.shape[-1]:
        raise ValueError(message)
    elif validate_args:
      assertions.append(
          tf.compat.v1.assert_equal(
              tf.shape(input=tensor_a)[-1],
              tf.shape(input=tensor_b)[-1],
              message=message))

  if logits is not None:
    validate_equal_last_dim(
        outcomes,
        logits,
        message='Last dimension of outcomes and logits must be equal size.')
  if probs is not None:
    validate_equal_last_dim(
        outcomes,
        probs,
        message='Last dimension of outcomes and probs must be equal size.')

  message = 'Rank of outcomes must be 1.'
  if outcomes.shape.ndims is not None:
    if outcomes.shape.ndims != 1:
      raise ValueError(message)
  elif validate_args:
    assertions.append(tf.compat.v1.assert_rank(outcomes, 1, message=message))

  message = 'Size of outcomes must be greater than 0.'
  if outcomes.shape.num_elements() is not None:
    if outcomes.shape.num_elements() == 0:
      raise ValueError(message)
  elif validate_args:
    assertions.append(
        tf.compat.v1.assert_greater(
            tf.size(input=outcomes), 0, message=message))

  if validate_args:
    assertions.append(
        tf.compat.v1.assert_equal(
            tf.math.is_strictly_increasing(outcomes),
            True,
            message='outcomes is not strictly increasing.'))

  return assertions

Validate outcomes logits and probs s shapes .


def _ensure_tf_install():  # pylint: disable=g-statement-before-imports
  """Attempt to import tensorflow, and ensure its version is sufficient.

  Raises:
    ImportError: if either tensorflow is not importable or its version is
    inadequate.
  """
  try:
    import tensorflow as tf
  except ImportError:
    # Print more informative error message, then reraise.
    print("\n\nFailed to import TensorFlow. Please note that TensorFlow is not "
          "installed by default when you install TensorFlow Probability. This "
          "is so that users can decide whether to install the GPU-enabled "
          "TensorFlow package. To use TensorFlow Probability, please install "
          "the most recent version of TensorFlow, by following instructions at "
          "https://tensorflow.org/install.\n\n")
    raise

  import distutils.version

  #
  # Update this whenever we need to depend on a newer TensorFlow release.
  #
  required_tensorflow_version = "1.13"

  if (distutils.version.LooseVersion(tf.__version__) <
      distutils.version.LooseVersion(required_tensorflow_version)):
    raise ImportError(
        "This version of TensorFlow Probability requires TensorFlow "
        "version >= {required}; Detected an installation of version {present}. "
        "Please upgrade TensorFlow to proceed.".format(
            required=required_tensorflow_version,
            present=tf.__version__))

Attempt to import tensorflow and ensure its version is sufficient .


def logistic_regression(features):
  """Bayesian logistic regression, which returns labels given features."""
  coeffs = ed.MultivariateNormalDiag(
      loc=tf.zeros(features.shape[1]), name="coeffs")
  labels = ed.Bernoulli(
      logits=tf.tensordot(features, coeffs, [[1], [0]]), name="labels")
  return labels

def logistic_regression ( features ) : coeffs = ed . MultivariateNormalDiag ( loc = tf . zeros ( features . shape [ 1 ] ) , name = "coeffs" ) labels = ed . Bernoulli ( logits = tf . tensordot ( features , coeffs , [ [ 1 ] , [ 0 ] ] ) , name = "labels" ) return labels

Bayesian logistic regression which returns labels given features .


def covertype():
  """Builds the Covertype data set."""
  import sklearn.datasets  # pylint: disable=g-import-not-at-top
  data = sklearn.datasets.covtype.fetch_covtype()
  features = data.data
  labels = data.target

  # Normalize features and append a column of ones for the intercept.
  features -= features.mean(0)
  features /= features.std(0)
  features = np.hstack([features, np.ones([features.shape[0], 1])])
  features = tf.cast(features, dtype=tf.float32)

  # Binarize outcomes on whether it is a specific category.
  _, counts = np.unique(labels, return_counts=True)
  specific_category = np.argmax(counts)
  labels = (labels == specific_category)
  labels = tf.cast(labels, dtype=tf.int32)
  return features, labels


def cholesky_covariance(x, sample_axis=0, keepdims=False, name=None):
  """Cholesky factor of the covariance matrix of vector-variate random samples.

  This function can be use to fit a multivariate normal to data.

  ```python
  tf.enable_eager_execution()
  import tensorflow_probability as tfp
  tfd = tfp.distributions

  # Assume data.shape = (1000, 2).  1000 samples of a random variable in R^2.
  observed_data = read_data_samples(...)

  # The mean is easy
  mu = tf.reduce_mean(observed_data, axis=0)

  # Get the scale matrix
  L = tfp.stats.cholesky_covariance(observed_data)

  # Make the best fit multivariate normal (under maximum likelihood condition).
  mvn = tfd.MultivariateNormalTriL(loc=mu, scale_tril=L)

  # Plot contours of the pdf.
  xs, ys = tf.meshgrid(
      tf.linspace(-5., 5., 50), tf.linspace(-5., 5., 50), indexing='ij')
  xy = tf.stack((tf.reshape(xs, [-1]), tf.reshape(ys, [-1])), axis=-1)
  pdf = tf.reshape(mvn.prob(xy), (50, 50))
  CS = plt.contour(xs, ys, pdf, 10)
  plt.clabel(CS, inline=1, fontsize=10)
  ```

  Why does this work?
  Given vector-variate random variables `X = (X1, ..., Xd)`, one may obtain the
  sample covariance matrix in `R^{d x d}` (see `tfp.stats.covariance`).

  The [Cholesky factor](https://en.wikipedia.org/wiki/Cholesky_decomposition)
  of this matrix is analogous to standard deviation for scalar random variables:
  Suppose `X` has covariance matrix `C`, with Cholesky factorization `C = L L^T`
  Then multiplying a vector of iid random variables which have unit variance by
  `L` produces a vector with covariance `L L^T`, which is the same as `X`.

  ```python
  observed_data = read_data_samples(...)
  L = tfp.stats.cholesky_covariance(observed_data, sample_axis=0)

  # Make fake_data with the same covariance as observed_data.
  uncorrelated_normal = tf.random_normal(shape=(500, 10))
  fake_data = tf.linalg.matvec(L, uncorrelated_normal)
  ```

  Args:
    x:  Numeric `Tensor`.  The rightmost dimension of `x` indexes events. E.g.
      dimensions of a random vector.
    sample_axis: Scalar or vector `Tensor` designating axis holding samples.
      Default value: `0` (leftmost dimension). Cannot be the rightmost dimension
        (since this indexes events).
    keepdims:  Boolean.  Whether to keep the sample axis as singletons.
    name: Python `str` name prefixed to Ops created by this function.
          Default value: `None` (i.e., `'covariance'`).

  Returns:
    chol:  `Tensor` of same `dtype` as `x`.  The last two dimensions hold
      lower triangular matrices (the Cholesky factors).
  """
  with tf.compat.v1.name_scope(
      name, 'cholesky_covariance', values=[x, sample_axis]):
    sample_axis = tf.convert_to_tensor(value=sample_axis, dtype=tf.int32)
    cov = covariance(
        x, sample_axis=sample_axis, event_axis=-1, keepdims=keepdims)
    return tf.linalg.cholesky(cov)

def cholesky_covariance ( x , sample_axis = 0 , keepdims = False , name = None ) : with tf . compat . v1 . name_scope ( name , 'cholesky_covariance' , values = [ x , sample_axis ] ) : sample_axis = tf . convert_to_tensor ( value = sample_axis , dtype = tf . int32 ) cov = covariance ( x , sample_axis = sample_axis , event_axis = - 1 , keepdims = keepdims ) return tf . linalg . cholesky ( cov )

Cholesky factor of the covariance matrix of vector - variate random samples .


def stddev(x, sample_axis=0, keepdims=False, name=None):
  """Estimate standard deviation using samples.

  Given `N` samples of scalar valued random variable `X`, standard deviation may
  be estimated as

  ```none
  Stddev[X] := Sqrt[Var[X]],
  Var[X] := N^{-1} sum_{n=1}^N (X_n - Xbar) Conj{(X_n - Xbar)},
  Xbar := N^{-1} sum_{n=1}^N X_n
  ```

  ```python
  x = tf.random_normal(shape=(100, 2, 3))

  # stddev[i, j] is the sample standard deviation of the (i, j) batch member.
  stddev = tfp.stats.stddev(x, sample_axis=0)
  ```

  Scaling a unit normal by a standard deviation produces normal samples
  with that standard deviation.

  ```python
  observed_data = read_data_samples(...)
  stddev = tfp.stats.stddev(observed_data)

  # Make fake_data with the same standard deviation as observed_data.
  fake_data = stddev * tf.random_normal(shape=(100,))
  ```

  Notice we divide by `N` (the numpy default), which does not create `NaN`
  when `N = 1`, but is slightly biased.

  Args:
    x:  A numeric `Tensor` holding samples.
    sample_axis: Scalar or vector `Tensor` designating axis holding samples, or
      `None` (meaning all axis hold samples).
      Default value: `0` (leftmost dimension).
    keepdims:  Boolean.  Whether to keep the sample axis as singletons.
    name: Python `str` name prefixed to Ops created by this function.
          Default value: `None` (i.e., `'stddev'`).

  Returns:
    stddev: A `Tensor` of same `dtype` as the `x`, and rank equal to
      `rank(x) - len(sample_axis)`
  """
  with tf.compat.v1.name_scope(name, 'stddev', values=[x, sample_axis]):
    return tf.sqrt(variance(x, sample_axis=sample_axis, keepdims=keepdims))

def stddev ( x , sample_axis = 0 , keepdims = False , name = None ) : with tf . compat . v1 . name_scope ( name , 'stddev' , values = [ x , sample_axis ] ) : return tf . sqrt ( variance ( x , sample_axis = sample_axis , keepdims = keepdims ) )

Estimate standard deviation using samples .


def variance(x, sample_axis=0, keepdims=False, name=None):
  """Estimate variance using samples.

  Given `N` samples of scalar valued random variable `X`, variance may
  be estimated as

  ```none
  Var[X] := N^{-1} sum_{n=1}^N (X_n - Xbar) Conj{(X_n - Xbar)}
  Xbar := N^{-1} sum_{n=1}^N X_n
  ```

  ```python
  x = tf.random_normal(shape=(100, 2, 3))

  # var[i, j] is the sample variance of the (i, j) batch member of x.
  var = tfp.stats.variance(x, sample_axis=0)
  ```

  Notice we divide by `N` (the numpy default), which does not create `NaN`
  when `N = 1`, but is slightly biased.

  Args:
    x:  A numeric `Tensor` holding samples.
    sample_axis: Scalar or vector `Tensor` designating axis holding samples, or
      `None` (meaning all axis hold samples).
      Default value: `0` (leftmost dimension).
    keepdims:  Boolean.  Whether to keep the sample axis as singletons.
    name: Python `str` name prefixed to Ops created by this function.
          Default value: `None` (i.e., `'variance'`).

  Returns:
    var: A `Tensor` of same `dtype` as the `x`, and rank equal to
      `rank(x) - len(sample_axis)`
  """
  with tf.compat.v1.name_scope(name, 'variance', values=[x, sample_axis]):
    return covariance(
        x, y=None, sample_axis=sample_axis, event_axis=None, keepdims=keepdims)

def variance ( x , sample_axis = 0 , keepdims = False , name = None ) : with tf . compat . v1 . name_scope ( name , 'variance' , values = [ x , sample_axis ] ) : return covariance ( x , y = None , sample_axis = sample_axis , event_axis = None , keepdims = keepdims )


def _make_positive_axis(axis, ndims):
  """Rectify possibly negatively axis. Prefer return Python list."""
  axis = _make_list_or_1d_tensor(axis)

  ndims = tf.convert_to_tensor(value=ndims, name='ndims', dtype=tf.int32)
  ndims_ = tf.get_static_value(ndims)

  if _is_list_like(axis) and ndims_ is not None:
    # Static case
    positive_axis = []
    for a in axis:
      if a < 0:
        a = ndims_ + a
      positive_axis.append(a)
  else:
    # Dynamic case
    axis = tf.convert_to_tensor(value=axis, name='axis', dtype=tf.int32)
    positive_axis = tf.where(axis >= 0, axis, axis + ndims)

  return positive_axis

def
_make_positive_axis
(
axis
,
ndims
)
:
axis
=
_make_list_or_1d_tensor
(
axis
)
ndims
=
tf
.
convert_to_tensor
(
value
=
ndims
,
name
=
'ndims'
,
dtype
=
tf
.
int32
)
ndims_
=
tf
.
get_static_value
(
ndims
)
if
_is_list_like
(
axis
)
and
ndims_
is
not
:
# Static case
positive_axis
=
[
]
for
a
in
axis
:
if
a
<
0
:
a
=
ndims_
+
a
positive_axis
.
append
(
a
)
else
:
# Dynamic case
axis
=
tf
.
convert_to_tensor
(
value
=
axis
,
name
=
'axis'
,
dtype
=
tf
.
int32
)
positive_axis
=
tf
.
where
(
axis
>=
0
,
axis
,
axis
+
ndims
)
return
positive_axis

Rectify possibly negatively axis . Prefer return Python list .


def _squeeze(x, axis):
  """A version of squeeze that works with dynamic axis."""
  x = tf.convert_to_tensor(value=x, name='x')
  if axis is None:
    return tf.squeeze(x, axis=None)
  axis = tf.convert_to_tensor(value=axis, name='axis', dtype=tf.int32)
  axis += tf.zeros([1], dtype=axis.dtype)  # Make axis at least 1d.
  keep_axis, _ = tf.compat.v1.setdiff1d(tf.range(0, tf.rank(x)), axis)
  return tf.reshape(x, tf.gather(tf.shape(input=x), keep_axis))

def
_squeeze
(
x
,
axis
)
:
x
=
tf
.
convert_to_tensor
(
value
=
x
,
name
=
'x'
)
if
axis
is
:
return
tf
.
squeeze
(
x
,
axis
=
)
axis
=
tf
.
convert_to_tensor
(
value
=
axis
,
name
=
'axis'
,
dtype
=
tf
.
int32
)
axis
+=
tf
.
zeros
(
[
1
]
,
dtype
=
axis
.
dtype
)
# Make axis at least 1d.
keep_axis
,
_
=
tf
.
compat
.
v1
.
setdiff1d
(
tf
.
range
(
0
,
tf
.
rank
(
x
)
)
,
axis
)
return
tf
.
reshape
(
x
,
tf
.
gather
(
tf
.
shape
(
input
=
x
)
,
keep_axis
)
)

A version of squeeze that works with dynamic axis .


def _z(self, x):
    """Standardize input `x` to a unit normal."""
    with tf.name_scope("standardize"):
      return (x - self.loc) / self.scale

def _z ( self , x ) : with tf . name_scope ( "standardize" ) : return ( x - self . loc ) / self . scale


def _inv_z(self, z):
    """Reconstruct input `x` from a its normalized version."""
    with tf.name_scope("reconstruct"):
      return z * self.scale + self.loc

def _inv_z ( self , z ) : with tf . name_scope ( "reconstruct" ) : return z * self . scale + self . loc

Reconstruct input x from a its normalized version .


def semilocal_linear_trend_transition_matrix(autoregressive_coef):
  """Build the transition matrix for a semi-local linear trend model."""
  # We want to write the following 2 x 2 matrix:
  #  [[1., 1., ],    # level(t+1) = level(t) + slope(t)
  #   [0., ar_coef], # slope(t+1) = ar_coef * slope(t)
  # but it's slightly tricky to properly incorporate the batch shape of
  # autoregressive_coef. E.g., if autoregressive_coef has shape [4,6], we want
  # to return shape [4, 6, 2, 2]. We do this by breaking the matrix into its
  # fixed entries, written explicitly, and then the autoregressive_coef part
  # which we add in after using a mask to broadcast to the correct matrix shape.

  fixed_entries = tf.constant(
      [[1., 1.],
       [0., 0.]],
      dtype=autoregressive_coef.dtype)

  autoregressive_coef_mask = tf.constant([[0., 0.],
                                          [0., 1.]],
                                         dtype=autoregressive_coef.dtype)
  bottom_right_entry = (autoregressive_coef[..., tf.newaxis, tf.newaxis] *
                        autoregressive_coef_mask)
  return tf.linalg.LinearOperatorFullMatrix(
      fixed_entries + bottom_right_entry)

Build the transition matrix for a semi - local linear trend model .


def semilocal_linear_trend_transition_noise(level_scale,
                                            slope_mean,
                                            slope_scale,
                                            autoregressive_coef):
  """Build the transition noise model for a semi-local linear trend model."""

  # At each timestep, the stochasticity of `level` and `slope` are given
  # by `level_scale` and `slope_scale` respectively.
  broadcast_batch_shape = dist_util.get_broadcast_shape(
      level_scale, slope_mean, slope_scale, autoregressive_coef)
  broadcast_ones = tf.ones(broadcast_batch_shape, dtype=level_scale.dtype)
  scale_diag = tf.stack([level_scale * broadcast_ones,
                         slope_scale * broadcast_ones],
                        axis=-1)

  # We additionally fold in a bias term implementing the nonzero `slope_mean`.
  # The overall `slope` update is (from `SemiLocalLinearTrend` docstring)
  #   slope[t] = (slope_mean +
  #               autoregressive_coef * (slope[t-1] - slope_mean) +
  #               Normal(0., slope_scale))
  # which we rewrite as
  #   slope[t] = (
  #    autoregressive_coef * slope[t-1] +                  # linear transition
  #    Normal(loc=slope_mean - autoregressive_coef * slope_mean,  # noise bias
  #           scale=slope_scale))                                 # noise scale
  bias = tf.stack([tf.zeros_like(broadcast_ones),
                   slope_mean * (1 - autoregressive_coef) * broadcast_ones],
                  axis=-1)
  return tfd.MultivariateNormalDiag(
      loc=bias,
      scale_diag=scale_diag)

#               Normal(0., slope_scale))

#           scale=slope_scale))                                 # noise scale

Build the transition noise model for a semi - local linear trend model .


def sample_halton_sequence(dim,
                           num_results=None,
                           sequence_indices=None,
                           dtype=tf.float32,
                           randomized=True,
                           seed=None,
                           name=None):
  r"""Returns a sample from the `dim` dimensional Halton sequence.

  Warning: The sequence elements take values only between 0 and 1. Care must be
  taken to appropriately transform the domain of a function if it differs from
  the unit cube before evaluating integrals using Halton samples. It is also
  important to remember that quasi-random numbers without randomization are not
  a replacement for pseudo-random numbers in every context. Quasi random numbers
  are completely deterministic and typically have significant negative
  autocorrelation unless randomization is used.

  Computes the members of the low discrepancy Halton sequence in dimension
  `dim`. The `dim`-dimensional sequence takes values in the unit hypercube in
  `dim` dimensions. Currently, only dimensions up to 1000 are supported. The
  prime base for the k-th axes is the k-th prime starting from 2. For example,
  if `dim` = 3, then the bases will be [2, 3, 5] respectively and the first
  element of the non-randomized sequence will be: [0.5, 0.333, 0.2]. For a more
  complete description of the Halton sequences see
  [here](https://en.wikipedia.org/wiki/Halton_sequence). For low discrepancy
  sequences and their applications see
  [here](https://en.wikipedia.org/wiki/Low-discrepancy_sequence).

  If `randomized` is true, this function produces a scrambled version of the
  Halton sequence introduced by [Owen (2017)][1]. For the advantages of
  randomization of low discrepancy sequences see [here](
  https://en.wikipedia.org/wiki/Quasi-Monte_Carlo_method#Randomization_of_quasi-Monte_Carlo).

  The number of samples produced is controlled by the `num_results` and
  `sequence_indices` parameters. The user must supply either `num_results` or
  `sequence_indices` but not both.
  The former is the number of samples to produce starting from the first
  element. If `sequence_indices` is given instead, the specified elements of
  the sequence are generated. For example, sequence_indices=tf.range(10) is
  equivalent to specifying n=10.

  #### Examples

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

  # Produce the first 1000 members of the Halton sequence in 3 dimensions.
  num_results = 1000
  dim = 3
  sample = tfp.mcmc.sample_halton_sequence(
    dim,
    num_results=num_results,
    seed=127)

  # Evaluate the integral of x_1 * x_2^2 * x_3^3  over the three dimensional
  # hypercube.
  powers = tf.range(1.0, limit=dim + 1)
  integral = tf.reduce_mean(tf.reduce_prod(sample ** powers, axis=-1))
  true_value = 1.0 / tf.reduce_prod(powers + 1.0)
  with tf.Session() as session:
    values = session.run((integral, true_value))

  # Produces a relative absolute error of 1.7%.
  print ("Estimated: %f, True Value: %f" % values)

  # Now skip the first 1000 samples and recompute the integral with the next
  # thousand samples. The sequence_indices argument can be used to do this.

  sequence_indices = tf.range(start=1000, limit=1000 + num_results,
                              dtype=tf.int32)
  sample_leaped = tfp.mcmc.sample_halton_sequence(
      dim,
      sequence_indices=sequence_indices,
      seed=111217)

  integral_leaped = tf.reduce_mean(tf.reduce_prod(sample_leaped ** powers,
                                                  axis=-1))
  with tf.Session() as session:
    values = session.run((integral_leaped, true_value))
  # Now produces a relative absolute error of 0.05%.
  print ("Leaped Estimated: %f, True Value: %f" % values)
  ```

  Args:
    dim: Positive Python `int` representing each sample's `event_size.` Must
      not be greater than 1000.
    num_results: (Optional) Positive scalar `Tensor` of dtype int32. The number
      of samples to generate. Either this parameter or sequence_indices must
      be specified but not both. If this parameter is None, then the behaviour
      is determined by the `sequence_indices`.
      Default value: `None`.
    sequence_indices: (Optional) `Tensor` of dtype int32 and rank 1. The
      elements of the sequence to compute specified by their position in the
      sequence. 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 None, then the `num_results` parameter must be specified
      which gives the number of desired samples starting from the first sample.
      Default value: `None`.
    dtype: (Optional) The dtype of the sample. One of: `float16`, `float32` or
      `float64`.
      Default value: `tf.float32`.
    randomized: (Optional) bool indicating whether to produce a randomized
      Halton sequence. If True, applies the randomization described in
      [Owen (2017)][1].
      Default value: `True`.
    seed: (Optional) Python integer to seed the random number generator. Only
      used if `randomized` is True. If not supplied and `randomized` is True,
      no seed is set.
      Default value: `None`.
    name:  (Optional) Python `str` describing ops managed by this function. If
      not supplied the name of this function is used.
      Default value: "sample_halton_sequence".

  Returns:
    halton_elements: Elements of the Halton sequence. `Tensor` of supplied dtype
      and `shape` `[num_results, dim]` if `num_results` was specified or shape
      `[s, dim]` where s is the size of `sequence_indices` if `sequence_indices`
      were specified.

  Raises:
    ValueError: if both `sequence_indices` and `num_results` were specified or
      if dimension `dim` is less than 1 or greater than 1000.

  #### References

  [1]: Art B. Owen. A randomized Halton algorithm in R. _arXiv preprint
       arXiv:1706.02808_, 2017. https://arxiv.org/abs/1706.02808
  """
  if dim < 1 or dim > _MAX_DIMENSION:
    raise ValueError(
        'Dimension must be between 1 and {}. Supplied {}'.format(_MAX_DIMENSION,
                                                                 dim))
  if (num_results is None) == (sequence_indices is None):
    raise ValueError('Either `num_results` or `sequence_indices` must be'
                     ' specified but not both.')

  if not dtype.is_floating:
    raise ValueError('dtype must be of `float`-type')

  with tf.compat.v1.name_scope(
      name, 'sample', values=[num_results, sequence_indices]):
    # Here and in the following, the shape layout is as follows:
    # [sample dimension, event dimension, coefficient dimension].
    # The coefficient dimension is an intermediate axes which will hold the
    # weights of the starting integer when expressed in the (prime) base for
    # an event dimension.
    if num_results is not None:
      num_results = tf.convert_to_tensor(value=num_results)
    if sequence_indices is not None:
      sequence_indices = tf.convert_to_tensor(value=sequence_indices)
    indices = _get_indices(num_results, sequence_indices, dtype)
    radixes = tf.constant(_PRIMES[0:dim], dtype=dtype, shape=[dim, 1])

    max_sizes_by_axes = _base_expansion_size(
        tf.reduce_max(input_tensor=indices), radixes)

    max_size = tf.reduce_max(input_tensor=max_sizes_by_axes)

    # The powers of the radixes that we will need. Note that there is a bit
    # of an excess here. Suppose we need the place value coefficients of 7
    # in base 2 and 3. For 2, we will have 3 digits but we only need 2 digits
    # for base 3. However, we can only create rectangular tensors so we
    # store both expansions in a [2, 3] tensor. This leads to the problem that
    # we might end up attempting to raise large numbers to large powers. For
    # example, base 2 expansion of 1024 has 10 digits. If we were in 10
    # dimensions, then the 10th prime (29) we will end up computing 29^10 even
    # though we don't need it. We avoid this by setting the exponents for each
    # axes to 0 beyond the maximum value needed for that dimension.
    exponents_by_axes = tf.tile([tf.range(max_size)], [dim, 1])

    # The mask is true for those coefficients that are irrelevant.
    weight_mask = exponents_by_axes >= max_sizes_by_axes
    capped_exponents = tf.where(
        weight_mask,
        tf.zeros_like(exponents_by_axes),
        exponents_by_axes)
    weights = radixes ** capped_exponents
    # The following computes the base b expansion of the indices. Suppose,
    # x = a0 + a1*b + a2*b^2 + ... Then, performing a floor div of x with
    # the vector (1, b, b^2, b^3, ...) will produce
    # (a0 + s1 * b, a1 + s2 * b, ...) where s_i are coefficients we don't care
    # about. Noting that all a_i < b by definition of place value expansion,
    # we see that taking the elements mod b of the above vector produces the
    # place value expansion coefficients.
    coeffs = tf.math.floordiv(indices, weights)
    coeffs *= 1. - tf.cast(weight_mask, dtype)
    coeffs %= radixes
    if not randomized:
      coeffs /= radixes
      return tf.reduce_sum(input_tensor=coeffs / weights, axis=-1)
    stream = distributions.SeedStream(seed, salt='MCMCSampleHaltonSequence')
    coeffs = _randomize(coeffs, radixes, seed=stream())
    # Remove the contribution from randomizing the trailing zero for the
    # axes where max_size_by_axes < max_size. This will be accounted
    # for separately below (using zero_correction).
    coeffs *= 1. - tf.cast(weight_mask, dtype)
    coeffs /= radixes
    base_values = tf.reduce_sum(input_tensor=coeffs / weights, axis=-1)

    # The randomization used in Owen (2017) does not leave 0 invariant. While
    # we have accounted for the randomization of the first `max_size_by_axes`
    # coefficients, we still need to correct for the trailing zeros. Luckily,
    # this is equivalent to adding a uniform random value scaled so the first
    # `max_size_by_axes` coefficients are zero. The following statements perform
    # this correction.
    zero_correction = tf.random.uniform([dim, 1], seed=stream(), dtype=dtype)
    zero_correction /= radixes ** max_sizes_by_axes
    return base_values + tf.reshape(zero_correction, [-1])

r Returns a sample from the dim dimensional Halton sequence .


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])

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]

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))

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)

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]

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
:
# 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
)

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

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__))

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)

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 .