def _joint_sample_n(self, n, seed=None):
    """Draw a joint sample from the prior over latents and observations."""

    with tf.name_scope("sample_n_joint"):
      stream = seed_stream.SeedStream(
          seed, salt="LinearGaussianStateSpaceModel_sample_n_joint")

      sample_and_batch_shape = distribution_util.prefer_static_value(
          tf.concat([[n], self.batch_shape_tensor()],
                    axis=0))

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

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

      initial_observation_matrix = (
          self.get_observation_matrix_for_timestep(self.initial_step))
      initial_observation_noise = (
          self.get_observation_noise_for_timestep(self.initial_step))

      initial_observation_pred = initial_observation_matrix.matmul(
          initial_latent)
      initial_observation = (initial_observation_pred +
                             initial_observation_noise.sample(
                                 sample_shape=_augment_sample_shape(
                                     initial_observation_noise,
                                     sample_and_batch_shape,
                                     self.validate_args),
                                 seed=stream())[..., tf.newaxis])

      sample_step = build_kalman_sample_step(
          self.get_transition_matrix_for_timestep,
          self.get_transition_noise_for_timestep,
          self.get_observation_matrix_for_timestep,
          self.get_observation_noise_for_timestep,
          full_sample_and_batch_shape=sample_and_batch_shape,
          stream=stream,
          validate_args=self.validate_args)

      # Scan over all timesteps to sample latents and observations.
      (latents, observations) = tf.scan(
          sample_step,
          elems=tf.range(self.initial_step+1, self.final_step),
          initializer=(initial_latent, initial_observation))

      # Combine the initial sampled timestep with the remaining timesteps.
      latents = tf.concat([initial_latent[tf.newaxis, ...],
                           latents], axis=0)
      observations = tf.concat([initial_observation[tf.newaxis, ...],
                                observations], axis=0)

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

    return latents, observations

Draw a joint sample from the prior over latents and observations .


def posterior_marginals(self, x, mask=None):
    """Run a Kalman smoother to return posterior mean and cov.

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

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

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

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

    with tf.name_scope("smooth"):
      x = tf.convert_to_tensor(value=x, name="x")
      (_, filtered_means, filtered_covs,
       predicted_means, predicted_covs, _, _) = self.forward_filter(
           x, mask=mask)

      (smoothed_means, smoothed_covs) = self.backward_smoothing_pass(
          filtered_means, filtered_covs,
          predicted_means, predicted_covs)

      return (smoothed_means, smoothed_covs)

def posterior_marginals ( self , x , mask = None ) : with tf . name_scope ( "smooth" ) : x = tf . convert_to_tensor ( value = x , name = "x" ) ( _ , filtered_means , filtered_covs , predicted_means , predicted_covs , _ , _ ) = self . forward_filter ( x , mask = mask ) ( smoothed_means , smoothed_covs ) = self . backward_smoothing_pass ( filtered_means , filtered_covs , predicted_means , predicted_covs ) return ( smoothed_means , smoothed_covs )

Run a Kalman smoother to return posterior mean and cov .


def _joint_mean(self):
    """Compute prior means for all variables via dynamic programming.

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

    with tf.name_scope("mean_joint"):

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

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

      initial_observation_mean = _propagate_mean(
          initial_latent_mean,
          self.get_observation_matrix_for_timestep(self.initial_step),
          self.get_observation_noise_for_timestep(self.initial_step))

      mean_step = build_kalman_mean_step(
          self.get_transition_matrix_for_timestep,
          self.get_transition_noise_for_timestep,
          self.get_observation_matrix_for_timestep,
          self.get_observation_noise_for_timestep)

      # Scan over all timesteps following the initial step.
      (latent_means, observation_means) = tf.scan(
          mean_step,
          elems=tf.range(self.initial_step+1, self.final_step),
          initializer=(initial_latent_mean, initial_observation_mean))

      # Squish the initial step back on top of the other (scanned) timesteps
      latent_means = tf.concat([initial_latent_mean[tf.newaxis, ...],
                                latent_means], axis=0)
      observation_means = tf.concat([initial_observation_mean[tf.newaxis, ...],
                                     observation_means], axis=0)

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

      return latent_means, observation_means

Compute prior means for all variables via dynamic programming .


def _joint_covariances(self):
    """Compute prior covariances for all variables via dynamic programming.

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

    with tf.name_scope("covariance_joint"):

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

      initial_observation_cov = _propagate_cov(
          initial_latent_cov,
          self.get_observation_matrix_for_timestep(self.initial_step),
          self.get_observation_noise_for_timestep(self.initial_step))

      cov_step = build_kalman_cov_step(
          self.get_transition_matrix_for_timestep,
          self.get_transition_noise_for_timestep,
          self.get_observation_matrix_for_timestep,
          self.get_observation_noise_for_timestep)

      # Scan over all timesteps following the initial step.
      (latent_covs, observation_covs) = tf.scan(
          cov_step,
          elems=tf.range(self.initial_step+1, self.final_step),
          initializer=(initial_latent_cov, initial_observation_cov))

      # Squish the initial step back on top of the other (scanned) timesteps
      latent_covs = tf.concat([initial_latent_cov[tf.newaxis, ...],
                               latent_covs], axis=0)
      observation_covs = tf.concat([initial_observation_cov[tf.newaxis, ...],
                                    observation_covs], axis=0)

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

Compute prior covariances for all variables via dynamic programming .


def latents_to_observations(self, latent_means, latent_covs):
    """Push latent means and covariances forward through the observation model.

    Args:
      latent_means: float `Tensor` of shape `[..., num_timesteps, latent_size]`
      latent_covs: float `Tensor` of shape
        `[..., num_timesteps, latent_size, latent_size]`.

    Returns:
      observation_means: float `Tensor` of shape
        `[..., num_timesteps, observation_size]`
      observation_covs: float `Tensor` of shape
        `[..., num_timesteps, observation_size, observation_size]`
    """

    with tf.name_scope("latents_to_observations"):

      pushforward_latents_step = build_pushforward_latents_step(
          self.get_observation_matrix_for_timestep,
          self.get_observation_noise_for_timestep)

      latent_means = distribution_util.move_dimension(
          latent_means, source_idx=-2, dest_idx=0)
      latent_means = latent_means[..., tf.newaxis]  # Make matmul happy.
      latent_covs = distribution_util.move_dimension(
          latent_covs, source_idx=-3, dest_idx=0)

      (initial_observation_mean,
       initial_observation_cov) = pushforward_latents_step(
           _=None,  # Loop body ignores previous observations.
           latent_t_mean_cov=(self.initial_step,
                              latent_means[self.initial_step],
                              latent_covs[self.initial_step]))

      # TODO(davmre) this loop is embarassingly parallel; replace with `pfor`.
      timesteps = tf.range(self.initial_step,
                           self.initial_step + self.num_timesteps)
      observation_means, observation_covs = tf.scan(
          pushforward_latents_step,
          elems=(timesteps, latent_means, latent_covs),
          initializer=(initial_observation_mean, initial_observation_cov),
          parallel_iterations=10000)

      observation_means = distribution_util.move_dimension(
          observation_means[..., 0], source_idx=0, dest_idx=-2)
      observation_covs = distribution_util.move_dimension(
          observation_covs, source_idx=0, dest_idx=-3)

      return observation_means, observation_covs

Push latent means and covariances forward through the observation model .


def _log_normalization(self):
    """Computes the log-normalizer of the distribution."""
    event_dim = tf.compat.dimension_value(self.event_shape[0])
    if event_dim is None:
      raise ValueError('vMF _log_normalizer currently only supports '
                       'statically known event shape')
    safe_conc = tf.where(self.concentration > 0,
                         self.concentration,
                         tf.ones_like(self.concentration))
    safe_lognorm = ((event_dim / 2 - 1) * tf.math.log(safe_conc) -
                    (event_dim / 2) * np.log(2 * np.pi) -
                    tf.math.log(_bessel_ive(event_dim / 2 - 1, safe_conc)) -
                    tf.abs(safe_conc))
    log_nsphere_surface_area = (
        np.log(2.) + (event_dim / 2) * np.log(np.pi) -
        tf.math.lgamma(tf.cast(event_dim / 2, self.dtype)))
    return tf.where(self.concentration > 0,
                    -safe_lognorm,
                    log_nsphere_surface_area * tf.ones_like(safe_lognorm))

Computes the log - normalizer of the distribution .


def _mode(self):
    """The mode of the von Mises-Fisher distribution is the mean direction."""
    return (self.mean_direction +
            tf.zeros_like(self.concentration)[..., tf.newaxis])

def _mode ( self ) : return ( self . mean_direction + tf . zeros_like ( self . concentration ) [ ... , tf . newaxis ] )

The mode of the von Mises - Fisher distribution is the mean direction .


def _rotate(self, samples):
    """Applies a Householder rotation to `samples`."""
    event_dim = (
        tf.compat.dimension_value(self.event_shape[0]) or
        self._event_shape_tensor()[0])
    basis = tf.concat([[1.], tf.zeros([event_dim - 1], dtype=self.dtype)],
                      axis=0),
    u = tf.nn.l2_normalize(basis - self.mean_direction, axis=-1)
    return samples - 2 * tf.reduce_sum(
        input_tensor=samples * u, axis=-1, keepdims=True) * u

def _rotate ( self , samples ) : event_dim = ( tf . compat . dimension_value ( self . event_shape [ 0 ] ) or self . _event_shape_tensor ( ) [ 0 ] ) basis = tf . concat ( [ [ 1. ] , tf . zeros ( [ event_dim - 1 ] , dtype = self . dtype ) ] , axis = 0 ) , u = tf . nn . l2_normalize ( basis - self . mean_direction , axis = - 1 ) return samples - 2 * tf . reduce_sum ( input_tensor = samples * u , axis = - 1 , keepdims = True ) * u

Applies a Householder rotation to samples .


def _sample_3d(self, n, seed=None):
    """Specialized inversion sampler for 3D."""
    seed = seed_stream.SeedStream(seed, salt='von_mises_fisher_3d')
    u_shape = tf.concat([[n], self._batch_shape_tensor()], axis=0)
    z = tf.random.uniform(u_shape, seed=seed(), dtype=self.dtype)
    # TODO(bjp): Higher-order odd dim analytic CDFs are available in [1], could
    # be bisected for bounded sampling runtime (i.e. not rejection sampling).
    # [1]: Inversion sampler via: https://ieeexplore.ieee.org/document/7347705/
    # The inversion is: u = 1 + log(z + (1-z)*exp(-2*kappa)) / kappa
    # We must protect against both kappa and z being zero.
    safe_conc = tf.where(self.concentration > 0,
                         self.concentration,
                         tf.ones_like(self.concentration))
    safe_z = tf.where(z > 0, z, tf.ones_like(z))
    safe_u = 1 + tf.reduce_logsumexp(
        input_tensor=[
            tf.math.log(safe_z),
            tf.math.log1p(-safe_z) - 2 * safe_conc
        ],
        axis=0) / safe_conc
    # Limit of the above expression as kappa->0 is 2*z-1
    u = tf.where(self.concentration > tf.zeros_like(safe_u), safe_u,
                 2 * z - 1)
    # Limit of the expression as z->0 is -1.
    u = tf.where(tf.equal(z, 0), -tf.ones_like(u), u)
    if not self._allow_nan_stats:
      u = tf.debugging.check_numerics(u, 'u in _sample_3d')
    return u[..., tf.newaxis]


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

  Args:
    fn: a callable

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

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


def _remove_dict_keys_with_value(dict_, val):
  """Removes `dict` keys which have have `self` as value."""
  return {k: v for k, v in dict_.items() if v is not val}

def _remove_dict_keys_with_value ( dict_ , val ) : return { k : v for k , v in dict_ . items ( ) if v is not val }

Removes dict keys which have have self as value .


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

Recursively replace dict s with _PrettyDict .


def _get_samples(dist, z, n, seed):
  """Check args and return samples."""
  with tf.compat.v1.name_scope('get_samples', values=[z, n]):
    if (n is None) == (z is None):
      raise ValueError(
          'Must specify exactly one of arguments "n" and "z".  Found: '
          'n = %s, z = %s' % (n, z))
    if n is not None:
      return dist.sample(n, seed=seed)
    else:
      return tf.convert_to_tensor(value=z, name='z')


def is_namedtuple_like(x):
  """Helper which returns `True` if input is `collections.namedtuple`-like."""
  try:
    for fn in x._fields:
      _ = getattr(x, fn)
    return True
  except AttributeError:
    return False

def is_namedtuple_like ( x ) : try : for fn in x . _fields : _ = getattr ( x , fn ) return True except AttributeError : return False

Helper which returns True if input is collections . namedtuple - like .


def _choose_base_case(is_accepted,
                      accepted,
                      rejected,
                      name=None):
  """Helper to `choose` which expand_dims `is_accepted` and applies tf.where."""
  def _expand_is_accepted_like(x):
    """Helper to expand `is_accepted` like the shape of some input arg."""
    with tf.compat.v1.name_scope('expand_is_accepted_like'):
      expand_shape = tf.concat([
          tf.shape(input=is_accepted),
          tf.ones([tf.rank(x) - tf.rank(is_accepted)], dtype=tf.int32),
      ],
                               axis=0)
      multiples = tf.concat([
          tf.ones([tf.rank(is_accepted)], dtype=tf.int32),
          tf.shape(input=x)[tf.rank(is_accepted):],
      ],
                            axis=0)
      m = tf.tile(tf.reshape(is_accepted, expand_shape),
                  multiples)
      m.set_shape(m.shape.merge_with(x.shape))
      return m
  def _where(accepted, rejected):
    if accepted is rejected:
      return accepted
    accepted = tf.convert_to_tensor(value=accepted, name='accepted')
    rejected = tf.convert_to_tensor(value=rejected, name='rejected')
    r = tf.where(_expand_is_accepted_like(accepted), accepted, rejected)
    r.set_shape(r.shape.merge_with(accepted.shape.merge_with(rejected.shape)))
    return r

  with tf.compat.v1.name_scope(
      name, 'choose', values=[is_accepted, accepted, rejected]):
    if not is_list_like(accepted):
      return _where(accepted, rejected)
    return [(choose(is_accepted, a, r, name=name) if is_namedtuple_like(a)
             else _where(a, r))
            for a, r in zip(accepted, rejected)]

Helper to choose which expand_dims is_accepted and applies tf . where .


def choose(is_accepted, accepted, rejected, name=None):
  """Helper which expand_dims `is_accepted` then applies tf.where."""
  if not is_namedtuple_like(accepted):
    return _choose_base_case(is_accepted, accepted, rejected, name=name)
  if not isinstance(accepted, type(rejected)):
    raise TypeError('Type of `accepted` ({}) must be identical to '
                    'type of `rejected` ({})'.format(
                        type(accepted).__name__,
                        type(rejected).__name__))
  return type(accepted)(**dict(
      [(fn,
        choose(is_accepted,
               getattr(accepted, fn),
               getattr(rejected, fn),
               name=name))
       for fn in accepted._fields]))

Helper which expand_dims is_accepted then applies tf . where .


def safe_sum(x, alt_value=-np.inf, name=None):
  """Elementwise adds list members, replacing non-finite results with alt_value.

  Typically the `alt_value` is chosen so the `MetropolisHastings`
  `TransitionKernel` always rejects the proposal.

  Args:
    x: Python `list` of `Tensors` to elementwise add.
    alt_value: Python scalar used to replace any elementwise sums which would
      otherwise be non-finite.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., "safe_sum").

  Returns:
    safe_sum: `Tensor` representing the elementwise sum of list of `Tensor`s
      `x` or `alt_value` where sums are non-finite.

  Raises:
    TypeError: if `x` is not list-like.
    ValueError: if `x` is empty.
  """
  with tf.compat.v1.name_scope(name, 'safe_sum', [x, alt_value]):
    if not is_list_like(x):
      raise TypeError('Expected list input.')
    if not x:
      raise ValueError('Input should not be empty.')
    in_shape = x[0].shape
    x = tf.stack(x, axis=-1)
    x = tf.reduce_sum(input_tensor=x, axis=-1)
    alt_value = np.array(alt_value, x.dtype.as_numpy_dtype)
    alt_fill = tf.fill(tf.shape(input=x), value=alt_value)
    x = tf.where(tf.math.is_finite(x), x, alt_fill)
    x.set_shape(x.shape.merge_with(in_shape))
    return x

def
safe_sum
(
x
,
alt_value
=
np
.
inf
,
name
=
)
:
with
tf
.
compat
.
v1
.
name_scope
(
name
,
'safe_sum'
,
[
x
,
alt_value
]
)
:
if
not
is_list_like
(
x
)
:
raise
TypeError
(
'Expected list input.'
)
if
not
x
:
raise
ValueError
(
'Input should not be empty.'
)
in_shape
=
x
[
0
]
.
shape
x
=
tf
.
stack
(
x
,
axis
=
1
)
x
=
tf
.
reduce_sum
(
input_tensor
=
x
,
axis
=
1
)
alt_value
=
np
.
array
(
alt_value
,
x
.
dtype
.
as_numpy_dtype
)
alt_fill
=
tf
.
fill
(
tf
.
shape
(
input
=
x
)
,
value
=
alt_value
)
x
=
tf
.
where
(
tf
.
math
.
is_finite
(
x
)
,
x
,
alt_fill
)
x
.
set_shape
(
x
.
shape
.
merge_with
(
in_shape
)
)
return
x

Elementwise adds list members replacing non - finite results with alt_value .


def _value_and_gradients(fn, fn_arg_list, result=None, grads=None, name=None):
  """Helper to `maybe_call_fn_and_grads`."""
  with tf.compat.v1.name_scope(name, 'value_and_gradients',
                               [fn_arg_list, result, grads]):

    def _convert_to_tensor(x, name):
      ctt = lambda x_: x_ if x_ is None else tf.convert_to_tensor(
          value=x_, name=name)
      return [ctt(x_) for x_ in x] if is_list_like(x) else ctt(x)

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

    if result is None:
      result = fn(*fn_arg_list)
      if grads is None and tf.executing_eagerly():
        # Ensure we disable bijector cacheing in eager mode.
        # TODO(b/72831017): Remove this once bijector cacheing is fixed for
        # eager mode.
        fn_arg_list = [0 + x for x in fn_arg_list]

    result = _convert_to_tensor(result, 'fn_result')

    if grads is not None:
      grads = _convert_to_tensor(grads, 'fn_grad')
      return result, grads

    if is_list_like(result) and len(result) == len(fn_arg_list):
      # Compute the block diagonal of Jacobian.
      # TODO(b/79158574): Guard this calculation by an arg which explicitly
      # requests block diagonal Jacobian calculation.
      def fn_slice(i):
        """Needed to prevent `cell-var-from-loop` pylint warning."""
        return lambda x: fn(*(fn_arg_list[:i] + [x] + fn_arg_list[i+1:]))
      grads = [
          tfp_math_value_and_gradients(fn_slice(i), fn_arg_list[i])[1]
          for i in range(len(result))
      ]
    else:
      _, grads = tfp_math_value_and_gradients(fn, fn_arg_list)

    return result, grads


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

Calls fn and computes the gradient of the result wrt args_list .


def smart_for_loop(loop_num_iter, body_fn, initial_loop_vars,
                   parallel_iterations=10, name=None):
  """Construct a for loop, preferring a python loop if `n` is staticaly known.

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

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

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

Construct a for loop preferring a python loop if n is staticaly known .


def trace_scan(loop_fn,
               initial_state,
               elems,
               trace_fn,
               parallel_iterations=10,
               name=None):
  """A simplified version of `tf.scan` that has configurable tracing.

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

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

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

    initial_state = tf.nest.map_structure(
        lambda x: tf.convert_to_tensor(value=x, name='initial_state'),
        initial_state)
    elems = tf.convert_to_tensor(value=elems, name='elems')

    static_length = elems.shape[0]
    if tf.compat.dimension_value(static_length) is None:
      length = tf.shape(input=elems)[0]
    else:
      length = tf.convert_to_tensor(
          value=static_length, dtype=tf.int32, name='length')

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

    trace_arrays = tf.nest.map_structure(
        lambda x: tf.TensorArray(x.dtype, size=length, element_shape=x.shape),
        trace_fn(initial_state))

    def _body(i, state, trace_arrays):
      state = loop_fn(state, elems_array.read(i))
      trace_arrays = tf.nest.pack_sequence_as(trace_arrays, [
          a.write(i, v) for a, v in zip(
              tf.nest.flatten(trace_arrays), tf.nest.flatten(trace_fn(state)))
      ])
      return i + 1, state, trace_arrays

    _, final_state, trace_arrays = tf.while_loop(
        cond=lambda i, *args: i < length,
        body=_body,
        loop_vars=(0, initial_state, trace_arrays),
        parallel_iterations=parallel_iterations)

    stacked_trace = tf.nest.map_structure(lambda x: x.stack(), trace_arrays)

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

    stacked_trace = tf.nest.map_structure(_merge_static_length, stacked_trace)
    return final_state, stacked_trace

A simplified version of tf . scan that has configurable tracing .


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

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

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

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

  @functools.wraps(setter)
  def _new_setter(kernel_results, *args, **kwargs):
    """Wrapped setter."""
    results_stack = []
    while hasattr(kernel_results, 'inner_results'):
      results_stack.append(kernel_results)
      kernel_results = kernel_results.inner_results

    new_kernel_results = setter(kernel_results, *args, **kwargs)
    for outer_results in reversed(results_stack):
      new_kernel_results = outer_results._replace(
          inner_results=new_kernel_results)

    return new_kernel_results

  return _new_setter

def
make_innermost_setter
(
setter
)
:
@
functools
.
wraps
(
setter
)
def
_new_setter
(
kernel_results
,
*
args
,
*
*
kwargs
)
:
"""Wrapped setter."""
results_stack
=
[
]
while
hasattr
(
kernel_results
,
'inner_results'
)
:
results_stack
.
append
(
kernel_results
)
kernel_results
=
kernel_results
.
inner_results
new_kernel_results
=
setter
(
kernel_results
,
*
args
,
*
*
kwargs
)
for
outer_results
in
reversed
(
results_stack
)
:
new_kernel_results
=
outer_results
.
_replace
(
inner_results
=
new_kernel_results
)
return
new_kernel_results
return
_new_setter

Wraps a setter so it applies to the inner - most results in kernel_results .


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

  The wrapped getter unwraps `kernel_results` and returns the return value of
  `getter` called with the first results without an `inner_results` attribute.

  Args:
    getter: A callable that takes Kernel results and returns some value.

  Returns:
    new_getter: A wrapped `getter`.
  """

  @functools.wraps(getter)
  def _new_getter(kernel_results, *args, **kwargs):
    """Wrapped getter."""
    results_stack = []
    while hasattr(kernel_results, 'inner_results'):
      results_stack.append(kernel_results)
      kernel_results = kernel_results.inner_results

    return getter(kernel_results, *args, **kwargs)

  return _new_getter

def
make_innermost_getter
(
getter
)
:
@
functools
.
wraps
(
getter
)
def
_new_getter
(
kernel_results
,
*
args
,
*
*
kwargs
)
:
"""Wrapped getter."""
results_stack
=
[
]
while
hasattr
(
kernel_results
,
'inner_results'
)
:
results_stack
.
append
(
kernel_results
)
kernel_results
=
kernel_results
.
inner_results
return
getter
(
kernel_results
,
*
args
,
*
*
kwargs
)
return
_new_getter

Wraps a getter so it applies to the inner - most results in kernel_results .


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

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

  Args:
    kernel: A TransitionKernel.

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

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

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

  for outer_kernel in reversed(kernel_stack):
    outer_kernel = _recreate_kernel(outer_kernel, {'inner_kernel': kernel})
    kernel = outer_kernel

  return kernel

def enable_store_parameters_in_results ( kernel ) : kernel_stack = [ ] while hasattr ( kernel , 'parameters' ) and 'inner_kernel' in kernel . parameters : kernel_stack . append ( kernel ) kernel = kernel . parameters [ 'inner_kernel' ] def _recreate_kernel ( kernel , parameters ) : new_parameters = kernel . parameters . copy ( ) new_parameters . update ( parameters ) if 'store_parameters_in_results' in new_parameters : new_parameters [ 'store_parameters_in_results' ] = True with deprecation . silence ( ) : return type ( kernel ) ( * * new_parameters ) if hasattr ( kernel , 'parameters' ) : kernel = _recreate_kernel ( kernel , { } ) for outer_kernel in reversed ( kernel_stack ) : outer_kernel = _recreate_kernel ( outer_kernel , { 'inner_kernel' : kernel } ) kernel = outer_kernel return kernel

Enables the store_parameters_in_results parameter in a chain of kernels .


def _replace_event_shape_in_shape_tensor(
    input_shape, event_shape_in, event_shape_out, validate_args):
  """Replaces the rightmost dims in a `Tensor` representing a shape.

  Args:
    input_shape: a rank-1 `Tensor` of integers
    event_shape_in: the event shape expected to be present in rightmost dims
      of `shape_in`.
    event_shape_out: the event shape with which to replace `event_shape_in` in
      the rightmost dims of `input_shape`.
    validate_args: Python `bool` indicating whether arguments should
      be checked for correctness.

  Returns:
    output_shape: A rank-1 integer `Tensor` with the same contents as
      `input_shape` except for the event dims, which are replaced with
      `event_shape_out`.
  """
  output_tensorshape, is_validated = _replace_event_shape_in_tensorshape(
      tensorshape_util.constant_value_as_shape(input_shape),
      event_shape_in,
      event_shape_out)

  # TODO(b/124240153): Remove map(tf.identity, deps) once tf.function
  # correctly supports control_dependencies.
  validation_dependencies = (
      map(tf.identity, (event_shape_in, event_shape_out))
      if validate_args else ())

  if (tensorshape_util.is_fully_defined(output_tensorshape) and
      (is_validated or not validate_args)):
    with tf.control_dependencies(validation_dependencies):
      output_shape = tf.convert_to_tensor(
          value=output_tensorshape, name='output_shape', dtype_hint=tf.int32)
    return output_shape, output_tensorshape

  with tf.control_dependencies(validation_dependencies):
    event_shape_in_ndims = (
        tf.size(input=event_shape_in)
        if tensorshape_util.num_elements(event_shape_in.shape) is None else
        tensorshape_util.num_elements(event_shape_in.shape))
    input_non_event_shape, input_event_shape = tf.split(
        input_shape, num_or_size_splits=[-1, event_shape_in_ndims])

  additional_assertions = []
  if is_validated:
    pass
  elif validate_args:
    # Check that `input_event_shape` and `event_shape_in` are compatible in the
    # sense that they have equal entries in any position that isn't a `-1` in
    # `event_shape_in`. Note that our validations at construction time ensure
    # there is at most one such entry in `event_shape_in`.
    mask = event_shape_in >= 0
    explicit_input_event_shape = tf.boolean_mask(
        tensor=input_event_shape, mask=mask)
    explicit_event_shape_in = tf.boolean_mask(
        tensor=event_shape_in, mask=mask)
    additional_assertions.append(
        assert_util.assert_equal(
            explicit_input_event_shape,
            explicit_event_shape_in,
            message='Input `event_shape` does not match `event_shape_in`.'))
    # We don't explicitly additionally verify
    # `tf.size(input_shape) > tf.size(event_shape_in)` since `tf.split`
    # already makes this assertion.

  with tf.control_dependencies(additional_assertions):
    output_shape = tf.concat([input_non_event_shape, event_shape_out], axis=0,
                             name='output_shape')

  return output_shape, output_tensorshape

Replaces the rightmost dims in a Tensor representing a shape .


def _replace_event_shape_in_tensorshape(
    input_tensorshape, event_shape_in, event_shape_out):
  """Replaces the event shape dims of a `TensorShape`.

  Args:
    input_tensorshape: a `TensorShape` instance in which to attempt replacing
      event shape.
    event_shape_in: `Tensor` shape representing the event shape expected to
      be present in (rightmost dims of) `tensorshape_in`. Must be compatible
      with the rightmost dims of `tensorshape_in`.
    event_shape_out: `Tensor` shape representing the new event shape, i.e.,
      the replacement of `event_shape_in`,

  Returns:
    output_tensorshape: `TensorShape` with the rightmost `event_shape_in`
      replaced by `event_shape_out`. Might be partially defined, i.e.,
      `TensorShape(None)`.
    is_validated: Python `bool` indicating static validation happened.

  Raises:
    ValueError: if we can determine the event shape portion of
      `tensorshape_in` as well as `event_shape_in` both statically, and they
      are not compatible. "Compatible" here means that they are identical on
      any dims that are not -1 in `event_shape_in`.
  """
  event_shape_in_ndims = tensorshape_util.num_elements(event_shape_in.shape)
  if tensorshape_util.rank(
      input_tensorshape) is None or event_shape_in_ndims is None:
    return tf.TensorShape(None), False  # Not is_validated.

  input_non_event_ndims = tensorshape_util.rank(
      input_tensorshape) - event_shape_in_ndims
  if input_non_event_ndims < 0:
    raise ValueError(
        'Input has fewer ndims ({}) than event shape ndims ({}).'.format(
            tensorshape_util.rank(input_tensorshape), event_shape_in_ndims))

  input_non_event_tensorshape = input_tensorshape[:input_non_event_ndims]
  input_event_tensorshape = input_tensorshape[input_non_event_ndims:]

  # Check that `input_event_shape_` and `event_shape_in` are compatible in the
  # sense that they have equal entries in any position that isn't a `-1` in
  # `event_shape_in`. Note that our validations at construction time ensure
  # there is at most one such entry in `event_shape_in`.
  event_shape_in_ = tf.get_static_value(event_shape_in)
  is_validated = (
      tensorshape_util.is_fully_defined(input_event_tensorshape) and
      event_shape_in_ is not None)
  if is_validated:
    input_event_shape_ = np.int32(input_event_tensorshape)
    mask = event_shape_in_ >= 0
    explicit_input_event_shape_ = input_event_shape_[mask]
    explicit_event_shape_in_ = event_shape_in_[mask]
    if not all(explicit_input_event_shape_ == explicit_event_shape_in_):
      raise ValueError(
          'Input `event_shape` does not match `event_shape_in`. '
          '({} vs {}).'.format(input_event_shape_, event_shape_in_))

  event_tensorshape_out = tensorshape_util.constant_value_as_shape(
      event_shape_out)
  if tensorshape_util.rank(event_tensorshape_out) is None:
    output_tensorshape = tf.TensorShape(None)
  else:
    output_tensorshape = tensorshape_util.concatenate(
        input_non_event_tensorshape, event_tensorshape_out)

  return output_tensorshape, is_validated

Replaces the event shape dims of a TensorShape .


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

  assertions = []

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

  shape_ = tf.get_static_value(shape)

  message = '`{}` elements must have at most one `-1`.'
  if shape_ is not None:
    if sum(shape_ == -1) > 1:
      raise ValueError(message.format(shape))
  elif validate_args:
    assertions.append(assert_util.assert_less(
        tf.reduce_sum(input_tensor=tf.cast(tf.equal(shape, -1), tf.int32)),
        2,
        message=message.format(shape)))

  message = '`{}` elements must be either positive integers or `-1`.'
  if shape_ is not None:
    if np.any(shape_ < -1):
      raise ValueError(message.format(shape))
  elif validate_args:
    assertions.append(assert_util.assert_greater(
        shape, -2, message=message.format(shape)))

  return assertions

Check that a shape Tensor is int - type and otherwise sane .


def converged_any(converged, failed):
  """Condition to stop when any batch member converges, or all have failed."""
  return (tf.reduce_any(input_tensor=converged) |
          tf.reduce_all(input_tensor=failed))

def converged_any ( converged , failed ) : return ( tf . reduce_any ( input_tensor = converged ) | tf . reduce_all ( input_tensor = failed ) )

Condition to stop when any batch member converges or all have failed .


def get_initial_state_args(value_and_gradients_function,
                           initial_position,
                           grad_tolerance,
                           control_inputs=None):
  """Returns a dictionary to populate the initial state of the search procedure.

  Performs an initial convergence check and the first evaluation of the
  objective function.

  Args:
    value_and_gradients_function: A Python callable that accepts a tensor and
      returns a tuple of two tensors: the objective function value and its
      derivative.
    initial_position: The starting point of the search procedure.
    grad_tolerance: The gradient tolerance for the procedure.
    control_inputs: Optional ops used to assert the validity of inputs, these
      are added as control dependencies to execute before the objective
      function is evaluated for the first time.

  Returns:
    An dictionary with values for the following keys:
      converged: True if the convergence check finds that the initial position
        is already an argmin of the objective function.
      failed: Initialized to False.
      num_objective_evaluations: Initialized to 1.
      position: Initialized to the initial position.
      objective_value: Initialized to the value of the objective function at
        the initial position.
      objective_gradient: Initialized to the gradient of the objective
        function at the initial position.
  """
  if control_inputs:
    with tf.control_dependencies(control_inputs):
      f0, df0 = value_and_gradients_function(initial_position)
  else:
    f0, df0 = value_and_gradients_function(initial_position)
  converged = norm(df0, dims=1) < grad_tolerance
  return dict(
      converged=converged,
      failed=tf.zeros_like(converged),  # i.e. False.
      num_iterations=tf.convert_to_tensor(value=0),
      num_objective_evaluations=tf.convert_to_tensor(value=1),
      position=initial_position,
      objective_value=f0,
      objective_gradient=df0)

def
get_initial_state_args
(
value_and_gradients_function
,
initial_position
,
grad_tolerance
,
control_inputs
=
)
:
if
control_inputs
:
with
tf
.
control_dependencies
(
control_inputs
)
:
f0
,
df0
=
value_and_gradients_function
(
initial_position
)
else
:
f0
,
df0
=
value_and_gradients_function
(
initial_position
)
converged
=
norm
(
df0
,
dims
=
1
)
<
grad_tolerance
return
dict
(
converged
=
converged
,
failed
=
tf
.
zeros_like
(
converged
)
,
# i.e. False.
num_iterations
=
tf
.
convert_to_tensor
(
value
=
0
)
,
num_objective_evaluations
=
tf
.
convert_to_tensor
(
value
=
1
)
,
position
=
initial_position
,
objective_value
=
f0
,
objective_gradient
=
df0
)

Returns a dictionary to populate the initial state of the search procedure .


def line_search_step(state, value_and_gradients_function, search_direction,
                     grad_tolerance, f_relative_tolerance, x_tolerance,
                     stopping_condition):
  """Performs the line search step of the BFGS search procedure.

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

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

  state_after_ls = update_fields(
      state,
      failed=state.failed | ~ls_result.converged,
      num_iterations=state.num_iterations + 1,
      num_objective_evaluations=(
          state.num_objective_evaluations + ls_result.func_evals))

  def _do_update_position():
    # For inactive batch members `left.x` is zero. However, their
    # `search_direction` might also be undefined, so we can't rely on
    # multiplication by zero to produce a `position_delta` of zero.
    position_delta = tf.where(
        inactive,
        tf.zeros_like(search_direction),
        search_direction * tf.expand_dims(ls_result.left.x, axis=-1))
    return _update_position(
        state_after_ls,
        position_delta,
        ls_result.left.f,
        ls_result.left.full_gradient,
        grad_tolerance, f_relative_tolerance, x_tolerance)

  return prefer_static.cond(
      stopping_condition(state.converged, state.failed),
      true_fn=lambda: state_after_ls,
      false_fn=_do_update_position)

Performs the line search step of the BFGS search procedure .


def _restrict_along_direction(value_and_gradients_function,
                              position,
                              direction):
  """Restricts a function in n-dimensions to a given direction.

  Suppose f: R^n -> R. Then given a point x0 and a vector p0 in R^n, the
  restriction of the function along that direction is defined by:

  ```None
  g(t) = f(x0 + t * p0)
  ```

  This function performs this restriction on the given function. In addition, it
  also computes the gradient of the restricted function along the restriction
  direction. This is equivalent to computing `dg/dt` in the definition above.

  Args:
    value_and_gradients_function: Callable accepting a single real `Tensor`
      argument of shape `[..., n]` and returning a tuple of a real `Tensor` of
      shape `[...]` and a real `Tensor` of shape `[..., n]`. The multivariate
      function whose restriction is to be computed. The output values of the
      callable are the function value and the gradients at the input argument.
    position: `Tensor` of real dtype and shape consumable by
      `value_and_gradients_function`. Corresponds to `x0` in the definition
      above.
    direction: `Tensor` of the same dtype and shape as `position`. The direction
      along which to restrict the function. Note that the direction need not
      be a unit vector.

  Returns:
    restricted_value_and_gradients_func: A callable accepting a tensor of shape
      broadcastable to `[...]` and same dtype as `position` and returning a
      namedtuple of `Tensors`. The input tensor is the parameter along the
      direction labelled `t` above. The return value contains fields:
        x: A real `Tensor` of shape `[...]`. The input value `t` where the line
          function was evaluated, after any necessary broadcasting.
        f: A real `Tensor` of shape `[...]` containing the value of the
          function at the point `position + t * direction`.
        df: A real `Tensor` of shape `[...]` containing the derivative at
          `position + t * direction`.
        full_gradient: A real `Tensor` of shape `[..., n]`, the full gradient
          of the original `value_and_gradients_function`.
  """
  def _restricted_func(t):
    t = _broadcast(t, position)
    pt = position + tf.expand_dims(t, axis=-1) * direction
    objective_value, gradient = value_and_gradients_function(pt)
    return ValueAndGradient(
        x=t,
        f=objective_value,
        df=tf.reduce_sum(input_tensor=gradient * direction, axis=-1),
        full_gradient=gradient)

  return _restricted_func

def _restrict_along_direction ( value_and_gradients_function , position , direction ) : def _restricted_func ( t ) : t = _broadcast ( t , position ) pt = position + tf . expand_dims ( t , axis = - 1 ) * direction objective_value , gradient = value_and_gradients_function ( pt ) return ValueAndGradient ( x = t , f = objective_value , df = tf . reduce_sum ( input_tensor = gradient * direction , axis = - 1 ) , full_gradient = gradient ) return _restricted_func

Restricts a function in n - dimensions to a given direction .


def _update_position(state,
                     position_delta,
                     next_objective,
                     next_gradient,
                     grad_tolerance,
                     f_relative_tolerance,
                     x_tolerance):
  """Updates the state advancing its position by a given position_delta."""
  failed = state.failed | ~tf.math.is_finite(next_objective) | ~tf.reduce_all(
      input_tensor=tf.math.is_finite(next_gradient), axis=-1)

  next_position = state.position + position_delta
  converged = ~failed & _check_convergence(state.position,
                                           next_position,
                                           state.objective_value,
                                           next_objective,
                                           next_gradient,
                                           grad_tolerance,
                                           f_relative_tolerance,
                                           x_tolerance)
  return update_fields(
      state,
      converged=state.converged | converged,
      failed=failed,
      position=next_position,
      objective_value=next_objective,
      objective_gradient=next_gradient)

def _update_position ( state , position_delta , next_objective , next_gradient , grad_tolerance , f_relative_tolerance , x_tolerance ) : failed = state . failed | ~ tf . math . is_finite ( next_objective ) | ~ tf . reduce_all ( input_tensor = tf . math . is_finite ( next_gradient ) , axis = - 1 ) next_position = state . position + position_delta converged = ~ failed & _check_convergence ( state . position , next_position , state . objective_value , next_objective , next_gradient , grad_tolerance , f_relative_tolerance , x_tolerance ) return update_fields ( state , converged = state . converged | converged , failed = failed , position = next_position , objective_value = next_objective , objective_gradient = next_gradient )

Updates the state advancing its position by a given position_delta .


def _check_convergence(current_position,
                       next_position,
                       current_objective,
                       next_objective,
                       next_gradient,
                       grad_tolerance,
                       f_relative_tolerance,
                       x_tolerance):
  """Checks if the algorithm satisfies the convergence criteria."""
  grad_converged = norm(next_gradient, dims=1) <= grad_tolerance
  x_converged = norm(next_position - current_position, dims=1) <= x_tolerance
  f_converged = (norm(next_objective - current_objective, dims=0) <=
                 f_relative_tolerance * current_objective)
  return grad_converged | x_converged | f_converged

Checks if the algorithm satisfies the convergence criteria .


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

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

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

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

def _broadcast ( value , target ) : return tf . broadcast_to ( tf . convert_to_tensor ( value = value , dtype = target . dtype ) , distribution_util . prefer_static_shape ( target ) [ : - 1 ] )

Broadcast a value to match the batching dimensions of a target .


def _harmonic_number(x):
  """Compute the harmonic number from its analytic continuation.

  Derivation from [here](
  https://en.wikipedia.org/wiki/Digamma_function#Relation_to_harmonic_numbers)
  and [Euler's constant](
  https://en.wikipedia.org/wiki/Euler%E2%80%93Mascheroni_constant).

  Args:
    x: input float.

  Returns:
    z: The analytic continuation of the harmonic number for the input.
  """
  one = tf.ones([], dtype=x.dtype)
  return tf.math.digamma(x + one) - tf.math.digamma(one)

def _harmonic_number ( x ) : one = tf . ones ( [ ] , dtype = x . dtype ) return tf . math . digamma ( x + one ) - tf . math . digamma ( one )

Compute the harmonic number from its analytic continuation .


def default_exchange_proposed_fn(prob_exchange):
  """Default exchange proposal function, for replica exchange MC.

  With probability `prob_exchange` propose combinations of replica for exchange.
  When exchanging, create combinations of adjacent replicas in
  [Replica Exchange Monte Carlo](
  https://en.wikipedia.org/wiki/Parallel_tempering)

  ```
  exchange_fn = default_exchange_proposed_fn(prob_exchange=0.5)
  exchange_proposed = exchange_fn(num_replica=3)

  exchange_proposed.eval()
  ==> [[0, 1]]  # 1 exchange, 0 <--> 1

  exchange_proposed.eval()
  ==> []  # 0 exchanges
  ```

  Args:
    prob_exchange: Scalar `Tensor` giving probability that any exchanges will
      be generated.

  Returns:
    default_exchange_proposed_fn_: Python callable which take a number of
      replicas (a Python integer), and return combinations of replicas for
      exchange as an [n, 2] integer `Tensor`, `0 <= n <= num_replica // 2`,
      with *unique* values in the set `{0, ..., num_replica}`.
  """

  def default_exchange_proposed_fn_(num_replica, seed=None):
    """Default function for `exchange_proposed_fn` of `kernel`."""
    seed_stream = distributions.SeedStream(seed, 'default_exchange_proposed_fn')

    zero_start = tf.random.uniform([], seed=seed_stream()) > 0.5
    if num_replica % 2 == 0:

      def _exchange():
        flat_exchange = tf.range(num_replica)
        if num_replica > 2:
          start = tf.cast(~zero_start, dtype=tf.int32)
          end = num_replica - start
          flat_exchange = flat_exchange[start:end]
        return tf.reshape(flat_exchange, [tf.size(input=flat_exchange) // 2, 2])
    else:

      def _exchange():
        start = tf.cast(zero_start, dtype=tf.int32)
        end = num_replica - tf.cast(~zero_start, dtype=tf.int32)
        flat_exchange = tf.range(num_replica)[start:end]
        return tf.reshape(flat_exchange, [tf.size(input=flat_exchange) // 2, 2])

    def _null_exchange():
      return tf.reshape(tf.cast([], dtype=tf.int32), shape=[0, 2])

    return tf.cond(
        pred=tf.random.uniform([], seed=seed_stream()) < prob_exchange,
        true_fn=_exchange,
        false_fn=_null_exchange)

  return default_exchange_proposed_fn_

Default exchange proposal function for replica exchange MC .


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

def
_get_field
(
kernel_results
,
field_name
)
:
if
hasattr
(
kernel_results
,
field_name
)
:
return
getattr
(
kernel_results
,
field_name
)
if
hasattr
(
kernel_results
,
'accepted_results'
)
:
return
getattr
(
kernel_results
.
accepted_results
,
field_name
)
raise
TypeError
(
'Cannot extract %s from %s'
%
(
field_name
,
kernel_results
)
)


def _get_exchanged_states(self, old_states, exchange_proposed,
                            exchange_proposed_n, sampled_replica_states,
                            sampled_replica_results):
    """Get list of TensorArrays holding exchanged states, and zeros."""
    with tf.compat.v1.name_scope('get_exchanged_states'):

      target_log_probs = []
      for replica in range(self.num_replica):
        replica_log_prob = _get_field(sampled_replica_results[replica],
                                      'target_log_prob')
        inverse_temp = self.inverse_temperatures[replica]
        target_log_probs.append(replica_log_prob / inverse_temp)
      target_log_probs = tf.stack(target_log_probs, axis=0)

      dtype = target_log_probs.dtype
      num_state_parts = len(sampled_replica_states[0])
      # exchanged_states[k][i] is Tensor of (new) state part k, for replica i.
      # The `k` will be known statically, and `i` is a Tensor.
      # We will insert values into indices `i` for every replica with a proposed
      # exchange.
      exchanged_states = [
          tf.TensorArray(
              dtype,
              size=self.num_replica,
              dynamic_size=False,
              tensor_array_name='exchanged_states',
              # State part k has same shape, regardless of replica.  So use 0.
              element_shape=sampled_replica_states[0][k].shape)
          for k in range(num_state_parts)
      ]

      # Draw random variables here, to avoid sampling in the loop (and losing
      # reproducibility).  This may mean we sample too many, but we will always
      # have enough.
      sample_shape = tf.concat(
          ([self.num_replica // 2], tf.shape(input=target_log_probs)[1:]),
          axis=0)
      log_uniforms = tf.math.log(
          tf.random.uniform(
              shape=sample_shape, dtype=dtype, seed=self._seed_stream()))

      def _swap(is_exchange_accepted, x, y):
        """Swap batches of x, y where accepted."""
        with tf.compat.v1.name_scope('swap_where_exchange_accepted'):
          new_x = mcmc_util.choose(is_exchange_accepted, y, x)
          new_y = mcmc_util.choose(is_exchange_accepted, x, y)
        return new_x, new_y

      def cond(i, unused_exchanged_states):
        return i < exchange_proposed_n

      def body(i, exchanged_states):
        """Body of while loop for exchanging states."""
        # Propose exchange between replicas indexed by m and n.
        m, n = tf.unstack(exchange_proposed[i])

        # Construct log_accept_ratio:  -temp_diff * target_log_prob_diff.
        # Note target_log_prob_diff = -EnergyDiff (common definition is in terms
        # of energy).
        temp_diff = self.inverse_temperatures[m] - self.inverse_temperatures[n]
        # Difference of target log probs may be +- Inf or NaN.  We want the
        # product of this with the temperature difference to have "alt value" of
        # -Inf.
        log_accept_ratio = mcmc_util.safe_sum(
            [-temp_diff * target_log_probs[m], temp_diff * target_log_probs[n]])

        is_exchange_accepted = log_uniforms[i] < log_accept_ratio

        for k in range(num_state_parts):
          new_m, new_n = _swap(is_exchange_accepted, old_states[k].read(m),
                               old_states[k].read(n))
          exchanged_states[k] = exchanged_states[k].write(m, new_m)
          exchanged_states[k] = exchanged_states[k].write(n, new_n)

        return i + 1, exchanged_states

      # At this point, exchanged_states[k] is a length num_replicas TensorArray.
      return tf.while_loop(
          cond=cond, body=body, loop_vars=[tf.constant(0),
                                           exchanged_states])[1]

Get list of TensorArrays holding exchanged states and zeros .


def _variance_scale_term(self):
    """Helper to `_covariance` and `_variance` which computes a shared scale."""
    # Expand back the last dim so the shape of _variance_scale_term matches the
    # shape of self.concentration.
    c0 = self.total_concentration[..., tf.newaxis]
    return tf.sqrt((1. + c0 / self.total_count[..., tf.newaxis]) / (1. + c0))

Helper to _covariance and _variance which computes a shared scale .


def forward_log_det_jacobian_fn(bijector):
  """Makes a function which applies a list of Bijectors' `log_det_jacobian`s."""
  if not mcmc_util.is_list_like(bijector):
    bijector = [bijector]

  def fn(transformed_state_parts, event_ndims):
    return sum([
        b.forward_log_det_jacobian(sp, event_ndims=e)
        for b, e, sp in zip(bijector, event_ndims, transformed_state_parts)
    ])

  return fn

def forward_log_det_jacobian_fn ( bijector ) : if not mcmc_util . is_list_like ( bijector ) : bijector = [ bijector ] def fn ( transformed_state_parts , event_ndims ) : return sum ( [ b . forward_log_det_jacobian ( sp , event_ndims = e ) for b , e , sp in zip ( bijector , event_ndims , transformed_state_parts ) ] ) return fn

Makes a function which applies a list of Bijectors log_det_jacobian s .


def forward_transform_fn(bijector):
  """Makes a function which applies a list of Bijectors' `forward`s."""
  if not mcmc_util.is_list_like(bijector):
    bijector = [bijector]

  def fn(transformed_state_parts):
    return [b.forward(sp) for b, sp in zip(bijector, transformed_state_parts)]

  return fn

def forward_transform_fn ( bijector ) : if not mcmc_util . is_list_like ( bijector ) : bijector = [ bijector ] def fn ( transformed_state_parts ) : return [ b . forward ( sp ) for b , sp in zip ( bijector , transformed_state_parts ) ] return fn

Makes a function which applies a list of Bijectors forward s .


def inverse_transform_fn(bijector):
  """Makes a function which applies a list of Bijectors' `inverse`s."""
  if not mcmc_util.is_list_like(bijector):
    bijector = [bijector]
  def fn(state_parts):
    return [b.inverse(sp)
            for b, sp in zip(bijector, state_parts)]
  return fn

def inverse_transform_fn ( bijector ) : if not mcmc_util . is_list_like ( bijector ) : bijector = [ bijector ] def fn ( state_parts ) : return [ b . inverse ( sp ) for b , sp in zip ( bijector , state_parts ) ] return fn

Makes a function which applies a list of Bijectors inverse s .


def one_step(self, current_state, previous_kernel_results):
    """Runs one iteration of the Transformed Kernel.

    Args:
      current_state: `Tensor` or Python `list` of `Tensor`s
        representing the current state(s) of the Markov chain(s),
        _after_ application of `bijector.forward`. The first `r`
        dimensions index independent chains,
        `r = tf.rank(target_log_prob_fn(*current_state))`. The
        `inner_kernel.one_step` does not actually use `current_state`,
        rather it takes as input
        `previous_kernel_results.transformed_state` (because
        `TransformedTransitionKernel` creates a copy of the input
        inner_kernel with a modified `target_log_prob_fn` which
        internally applies the `bijector.forward`).
      previous_kernel_results: `collections.namedtuple` containing `Tensor`s
        representing values from previous calls to this function (or from the
        `bootstrap_results` function.)

    Returns:
      next_state: Tensor or Python list of `Tensor`s representing the state(s)
        of the Markov chain(s) after taking exactly one step. Has same type and
        shape as `current_state`.
      kernel_results: `collections.namedtuple` of internal calculations used to
        advance the chain.
    """
    with tf.compat.v1.name_scope(
        name=mcmc_util.make_name(self.name, 'transformed_kernel', 'one_step'),
        values=[previous_kernel_results]):
      transformed_next_state, kernel_results = self._inner_kernel.one_step(
          previous_kernel_results.transformed_state,
          previous_kernel_results.inner_results)
      transformed_next_state_parts = (
          transformed_next_state
          if mcmc_util.is_list_like(transformed_next_state) else
          [transformed_next_state])
      next_state_parts = self._forward_transform(transformed_next_state_parts)
      next_state = (
          next_state_parts if mcmc_util.is_list_like(transformed_next_state)
          else next_state_parts[0])
      kernel_results = TransformedTransitionKernelResults(
          transformed_state=transformed_next_state,
          inner_results=kernel_results)
      return next_state, kernel_results

Runs one iteration of the Transformed Kernel .


def val_where(cond, tval, fval):
  """Like tf.where but works on namedtuples."""
  if isinstance(tval, tf.Tensor):
    return tf.where(cond, tval, fval)
  elif isinstance(tval, tuple):
    cls = type(tval)
    return cls(*(val_where(cond, t, f) for t, f in zip(tval, fval)))
  else:
    raise Exception(TypeError)

def val_where ( cond , tval , fval ) : if isinstance ( tval , tf . Tensor ) : return tf . where ( cond , tval , fval ) elif isinstance ( tval , tuple ) : cls = type ( tval ) return cls ( * ( val_where ( cond , t , f ) for t , f in zip ( tval , fval ) ) ) else : raise Exception ( TypeError )

Like tf . where but works on namedtuples .


def secant2(value_and_gradients_function,
            val_0,
            search_interval,
            f_lim,
            sufficient_decrease_param=0.1,
            curvature_param=0.9,
            name=None):
  """Performs the secant square procedure of Hager Zhang.

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

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

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

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

    initial_args = _Secant2Result(
        active=~failed & ~converged,
        converged=converged,
        failed=failed,
        num_evals=search_interval.func_evals + 1,
        left=val_left,
        right=val_right)

    def _apply_secant2_inner():
      return _secant2_inner(
          value_and_gradients_function,
          initial_args,
          val_0,
          val_c,
          f_lim,
          sufficient_decrease_param,
          curvature_param)

    return prefer_static.cond(
        tf.reduce_any(input_tensor=initial_args.active),
        _apply_secant2_inner,
        lambda: initial_args)

Performs the secant square procedure of Hager Zhang .


def _secant2_inner(value_and_gradients_function,
                   initial_args,
                   val_0,
                   val_c,
                   f_lim,
                   sufficient_decrease_param,
                   curvature_param):
  """Helper function for secant square."""
  # Apply the `update` function on active branch members to squeeze their
  # bracketing interval.
  update_result = update(value_and_gradients_function,
                         initial_args.left,
                         initial_args.right,
                         val_c,
                         f_lim,
                         active=initial_args.active)

  # Update active and failed flags, update left/right on non-failed entries.
  active = initial_args.active & ~update_result.failed
  failed = initial_args.failed | update_result.failed
  val_left = val_where(active, update_result.left, initial_args.left)
  val_right = val_where(active, update_result.right, initial_args.right)

  # Check if new `c` points should be generated.
  updated_left = active & tf.equal(val_left.x, val_c.x)
  updated_right = active & tf.equal(val_right.x, val_c.x)
  is_new = updated_left | updated_right

  next_c = tf.where(updated_left,
                    _secant(initial_args.left, val_left),
                    val_c.x)
  next_c = tf.where(updated_right,
                    _secant(initial_args.right, val_right),
                    next_c)
  in_range = (val_left.x <= next_c) & (next_c <= val_right.x)

  # Figure out if an extra function evaluation is needed for new `c` points.
  needs_extra_eval = tf.reduce_any(input_tensor=in_range & is_new)
  num_evals = initial_args.num_evals + update_result.num_evals
  num_evals = num_evals + tf.cast(needs_extra_eval, num_evals.dtype)

  next_args = _Secant2Result(
      active=active & in_range,  # No longer active if `c` is out of range.
      converged=initial_args.converged,
      failed=failed,
      num_evals=num_evals,
      left=val_left,
      right=val_right)

  def _apply_inner_update():
    next_val_c = prefer_static.cond(
        needs_extra_eval,
        (lambda: value_and_gradients_function(next_c)),
        (lambda: val_c))
    return _secant2_inner_update(
        value_and_gradients_function, next_args, val_0, next_val_c, f_lim,
        sufficient_decrease_param, curvature_param)

  return prefer_static.cond(
      tf.reduce_any(input_tensor=next_args.active),
      _apply_inner_update,
      lambda: next_args)


def _secant2_inner_update(value_and_gradients_function,
                          initial_args,
                          val_0,
                          val_c,
                          f_lim,
                          sufficient_decrease_param,
                          curvature_param):
  """Helper function for secant-square step."""
  # Fail if `val_c` is no longer finite.
  new_failed = initial_args.active & ~is_finite(val_c)
  active = initial_args.active & ~new_failed
  failed = initial_args.failed | new_failed

  # We converge when we find a point satisfying the Wolfe conditions, in those
  # cases we set `val_left = val_right = val_c`.
  found_wolfe = active & _satisfies_wolfe(
      val_0, val_c, f_lim, sufficient_decrease_param, curvature_param)
  val_left = val_where(found_wolfe, val_c, initial_args.left)
  val_right = val_where(found_wolfe, val_c, initial_args.right)
  converged = initial_args.converged | found_wolfe
  active = active & ~found_wolfe

  # If any active batch members remain, we apply the `update` function to
  # squeeze further their corresponding left/right bracketing interval.
  def _apply_update():
    update_result = update(
        value_and_gradients_function, val_left, val_right, val_c, f_lim,
        active=active)
    return _Secant2Result(
        active=tf.zeros_like(active),  # End of secant2, no actives anymore.
        converged=converged,
        failed=failed | update_result.failed,
        num_evals=initial_args.num_evals + update_result.num_evals,
        left=update_result.left,
        right=update_result.right)

  # Otherwise just return the current results.
  def _default():
    return _Secant2Result(
        active=active,
        converged=converged,
        failed=failed,
        num_evals=initial_args.num_evals,
        left=val_left,
        right=val_right)

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

Helper function for secant - square step .


def update(value_and_gradients_function, val_left, val_right, val_trial, f_lim,
           active=None):
  """Squeezes a bracketing interval containing the minimum.

  Given an interval which brackets a minimum and a point in that interval,
  finds a smaller nested interval which also brackets the minimum. If the
  supplied point does not lie in the bracketing interval, the current interval
  is returned.

  The following description is given in terms of individual points evaluated on
  a line function to be minimized. Note, however, the implementation also
  accepts batches of points allowing to minimize multiple line functions at
  once. See details on the docstring of `value_and_gradients_function` below.

  The requirement of the interval bracketing a minimum is expressed through the
  opposite slope conditions. Assume the left end point is 'a', the right
  end point is 'b', the function to be minimized is 'f' and the derivative is
  'df'. The update procedure relies on the following conditions being satisfied:

  '''
    f(a) <= f(0) + epsilon   (1)
    df(a) < 0                (2)
    df(b) > 0                (3)
  '''

  In the first condition, epsilon is a small positive constant. The condition
  demands that the function at the left end point be not much bigger than the
  starting point (i.e. 0). This is an easy to satisfy condition because by
  assumption, we are in a direction where the function value is decreasing.
  The second and third conditions together demand that there is at least one
  zero of the derivative in between a and b.

  In addition to the interval, the update algorithm requires a third point to
  be supplied. Usually, this point would lie within the interval [a, b]. If the
  point is outside this interval, the current interval is returned. If the
  point lies within the interval, the behaviour of the function and derivative
  value at this point is used to squeeze the original interval in a manner that
  preserves the opposite slope conditions.

  For further details of this component, see the procedure U0-U3 on page 123 of
  the [Hager and Zhang (2006)][2] article.

  Note that this function does not explicitly verify whether the opposite slope
  conditions are satisfied for the supplied interval. It is assumed that this
  is so.

  Args:
    value_and_gradients_function: A Python callable that accepts a real scalar
      tensor and returns an object that can be converted to a namedtuple.
      The namedtuple should have fields 'f' and 'df' that correspond to scalar
      tensors of real dtype containing the value of the function and its
      derivative at that point. The other namedtuple fields, if present,
      should be tensors or sequences (possibly nested) of tensors.
      In usual optimization application, this function would be generated by
      projecting the multivariate objective function along some specific
      direction. The direction is determined by some other procedure but should
      be a descent direction (i.e. the derivative of the projected univariate
      function must be negative at 0.).
      Alternatively, the function may represent the batching of `n` such line
      functions (e.g. projecting a single multivariate objective function along
      `n` distinct directions at once) accepting n points as input, i.e. a
      tensor of shape [n], and the fields 'f' and 'df' in the returned
      namedtuple should each be a tensor of shape [n], with the corresponding
      function values and derivatives at the input points.
    val_left: Return value of value_and_gradients_function at the left
      end point of the bracketing interval (labelles 'a' above).
    val_right: Return value of value_and_gradients_function at the right
      end point of the bracketing interval (labelles 'b' above).
    val_trial: Return value of value_and_gradients_function at the trial point
      to be used to shrink the interval (labelled 'c' above).
    f_lim: real `Tensor` of shape [n]. The function value threshold for
      the approximate Wolfe conditions to be checked for each batch member.
    active: optional boolean `Tensor` of shape [n]. Relevant in batching mode
      only, indicates batch members on which the update procedure should be
      applied. On non-active members the current left/right interval is returned
      unmodified.

  Returns:
    A namedtuple containing the following fields:
      iteration: An int32 scalar `Tensor`. The number of iterations performed
        by the bisect algorithm.
      stopped: A boolean `Tensor` of shape [n]. True for those batch members
        where the bisection algorithm terminated.
      failed: A boolean `Tensor` of shape [n]. True for those batch members
        where an error was encountered.
      num_evals: An int32 scalar `Tensor`. The number of times the objective
        function was evaluated.
      left: Return value of value_and_gradients_function at the updated left
        end point of the interval found.
      right: Return value of value_and_gradients_function at the updated right
        end point of the interval found.
  """
  # We should only update if the trial point is within the interval.
  within_range = (val_left.x < val_trial.x) & (val_trial.x < val_right.x)
  if active is not None:
    within_range = within_range & active

  # The new point is a valid left end point if it has negative slope
  # and the value at the point is not too large.
  valid_left = (val_trial.df < 0) & (val_trial.f <= f_lim)

  # If the trial point has a negative slope but the value at that point
  # is too high, bisect can narrow down an interval between the current left
  # and the trial point.
  needs_bisect = within_range & (val_trial.df < 0) & (val_trial.f > f_lim)

  # Note that if `~valid_left` it is because either:
  # - the slope at the trial point is positive, so it is a valid right
  #   point, or
  # - the needs_bisect condition is true.
  # In both cases we want to keep the current left and replace right
  # with the trial point.
  left = val_where(within_range & valid_left, val_trial, val_left)
  right = val_where(within_range & ~valid_left, val_trial, val_right)

  bisect_args = _IntermediateResult(
      iteration=tf.convert_to_tensor(value=0),
      stopped=~needs_bisect,
      failed=tf.zeros_like(within_range),  # i.e. all false.
      num_evals=tf.convert_to_tensor(value=0),
      left=left,
      right=right)
  return _bisect(value_and_gradients_function, bisect_args, f_lim)

Squeezes a bracketing interval containing the minimum .


def bracket(value_and_gradients_function,
            search_interval,
            f_lim,
            max_iterations,
            expansion_param=5.0):
  """Brackets the minimum given an initial starting point.

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

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

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

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

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

  # In these three cases bracketing is already `stopped` and there is no need
  # to perform further evaluations. Otherwise the bracketing loop is needed to
  # expand the interval, step B3, until the conditions are met.
  initial_args = _IntermediateResult(
      iteration=search_interval.iterations,
      stopped=already_stopped | bracketed | needs_bisect,
      failed=search_interval.failed,
      num_evals=search_interval.func_evals,
      left=search_interval.left,
      right=search_interval.right)

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

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

    # Updated the failed, bracketed, and needs_bisect conditions.
    failed = curr.failed | ~is_finite(right)
    bracketed = right.df >= 0
    needs_bisect = (right.df < 0) & (right.f > f_lim)
    return [_IntermediateResult(
        iteration=curr.iteration + 1,
        stopped=curr.stopped | failed | bracketed | needs_bisect,
        failed=failed,
        num_evals=curr.num_evals + 1,
        left=left,
        right=right)]

  bracket_result = tf.while_loop(
      cond=_loop_cond, body=_loop_body, loop_vars=[initial_args])[0]

  # For entries where bisect is still needed, mark them as not yet stopped,
  # reset the left end point, and run `_bisect` on them.
  needs_bisect = (
      (bracket_result.right.df < 0) & (bracket_result.right.f > f_lim))
  stopped = already_stopped | bracket_result.failed | ~needs_bisect
  left = val_where(stopped, bracket_result.left, search_interval.left)
  bisect_args = bracket_result._replace(stopped=stopped, left=left)
  return _bisect(value_and_gradients_function, bisect_args, f_lim)

Brackets the minimum given an initial starting point .