def _move_dims_to_flat_end(x, axis, x_ndims, right_end=True):
  """Move dims corresponding to `axis` in `x` to the end, then flatten.

  Args:
    x: `Tensor` with shape `[B0,B1,...,Bb]`.
    axis:  Python list of indices into dimensions of `x`.
    x_ndims:  Python integer holding number of dimensions in `x`.
    right_end:  Python bool.  Whether to move dims to the right end (else left).

  Returns:
    `Tensor` with value from `x` and dims in `axis` moved to end into one single
      dimension.
  """

  if not axis:
    return x

  # Suppose x.shape = [a, b, c, d]
  # Suppose axis = [1, 3]

  # other_dims = [0, 2] in example above.
  other_dims = sorted(set(range(x_ndims)).difference(axis))
  # x_permed.shape = [a, c, b, d]
  perm = other_dims + list(axis) if right_end else list(axis) + other_dims
  x_permed = tf.transpose(a=x, perm=perm)

  if x.shape.is_fully_defined():
    x_shape = x.shape.as_list()
    # other_shape = [a, c], end_shape = [b * d]
    other_shape = [x_shape[i] for i in other_dims]
    end_shape = [np.prod([x_shape[i] for i in axis])]
    full_shape = (
        other_shape + end_shape if right_end else end_shape + other_shape)
  else:
    other_shape = tf.gather(tf.shape(input=x), other_dims)
    full_shape = tf.concat(
        [other_shape, [-1]] if right_end else [[-1], other_shape], axis=0)
  return tf.reshape(x_permed, shape=full_shape)

Move dims corresponding to axis in x to the end then flatten .


def _sort_tensor(tensor):
  """Use `top_k` to sort a `Tensor` along the last dimension."""
  sorted_, _ = tf.nn.top_k(tensor, k=tf.shape(input=tensor)[-1])
  sorted_.set_shape(tensor.shape)
  return sorted_

def _sort_tensor ( tensor ) : sorted_ , _ = tf . nn . top_k ( tensor , k = tf . shape ( input = tensor ) [ - 1 ] ) sorted_ . set_shape ( tensor . shape ) return sorted_

Use top_k to sort a Tensor along the last dimension .


def make_component_state_space_models(self,
                                        num_timesteps,
                                        param_vals,
                                        initial_step=0):
    """Build an ordered list of Distribution instances for component models.

    Args:
      num_timesteps: Python `int` number of timesteps to model.
      param_vals: a list of `Tensor` parameter values in order corresponding to
        `self.parameters`, or a dict mapping from parameter names to values.
      initial_step: optional `int` specifying the initial timestep to model.
        This is relevant when the model contains time-varying components,
        e.g., holidays or seasonality.

    Returns:
      component_ssms: a Python list of `LinearGaussianStateSpaceModel`
        Distribution objects, in order corresponding to `self.components`.
    """

    with tf.compat.v1.name_scope('make_component_state_space_models'):

      # List the model parameters in canonical order
      param_map = self._canonicalize_param_vals_as_map(param_vals)
      param_vals_list = [param_map[p.name] for p in self.parameters]

      # Build SSMs for each component model. We process the components in
      # canonical order, extracting the parameters for each component from the
      # (ordered) list of parameters.
      remaining_param_vals = param_vals_list[1:]
      component_ssms = []
      for component in self.components:
        num_parameters = len(component.parameters)
        component_param_vals = remaining_param_vals[:num_parameters]
        remaining_param_vals = remaining_param_vals[num_parameters:]

        component_ssms.append(
            component.make_state_space_model(
                num_timesteps,
                param_vals=component_param_vals,
                initial_step=initial_step))

    return component_ssms

Build an ordered list of Distribution instances for component models .


def amari_alpha(logu, alpha=1., self_normalized=False, name=None):
  """The Amari-alpha Csiszar-function in log-space.

  A Csiszar-function is a member of,

  ```none
  F = { f:R_+ to R : f convex }.
  ```

  When `self_normalized = True`, the Amari-alpha Csiszar-function is:

  ```none
  f(u) = { -log(u) + (u - 1),     alpha = 0
         { u log(u) - (u - 1),    alpha = 1
         { [(u**alpha - 1) - alpha (u - 1)] / (alpha (alpha - 1)),    otherwise
  ```

  When `self_normalized = False` the `(u - 1)` terms are omitted.

  Warning: when `alpha != 0` and/or `self_normalized = True` this function makes
  non-log-space calculations and may therefore be numerically unstable for
  `|logu| >> 0`.

  For more information, see:
    A. Cichocki and S. Amari. "Families of Alpha-Beta-and GammaDivergences:
    Flexible and Robust Measures of Similarities." Entropy, vol. 12, no. 6, pp.
    1532-1568, 2010.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    alpha: `float`-like Python scalar. (See Mathematical Details for meaning.)
    self_normalized: Python `bool` indicating whether `f'(u=1)=0`. When
      `f'(u=1)=0` the implied Csiszar f-Divergence remains non-negative even
      when `p, q` are unnormalized measures.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    amari_alpha_of_u: `float`-like `Tensor` of the Csiszar-function evaluated
      at `u = exp(logu)`.

  Raises:
    TypeError: if `alpha` is `None` or a `Tensor`.
    TypeError: if `self_normalized` is `None` or a `Tensor`.
  """
  with tf.compat.v1.name_scope(name, "amari_alpha", [logu]):
    if alpha is None or tf.is_tensor(alpha):
      raise TypeError("`alpha` cannot be `None` or `Tensor` type.")
    if (self_normalized is None or tf.is_tensor(self_normalized)):
      raise TypeError("`self_normalized` cannot be `None` or `Tensor` type.")

    logu = tf.convert_to_tensor(value=logu, name="logu")

    if alpha == 0.:
      f = -logu
    elif alpha == 1.:
      f = tf.exp(logu) * logu
    else:
      f = tf.math.expm1(alpha * logu) / (alpha * (alpha - 1.))

    if not self_normalized:
      return f

    if alpha == 0.:
      return f + tf.math.expm1(logu)
    elif alpha == 1.:
      return f - tf.math.expm1(logu)
    else:
      return f - tf.math.expm1(logu) / (alpha - 1.)

The Amari - alpha Csiszar - function in log - space .


def kl_reverse(logu, self_normalized=False, name=None):
  """The reverse Kullback-Leibler Csiszar-function in log-space.

  A Csiszar-function is a member of,

  ```none
  F = { f:R_+ to R : f convex }.
  ```

  When `self_normalized = True`, the KL-reverse Csiszar-function is:

  ```none
  f(u) = -log(u) + (u - 1)
  ```

  When `self_normalized = False` the `(u - 1)` term is omitted.

  Observe that as an f-Divergence, this Csiszar-function implies:

  ```none
  D_f[p, q] = KL[q, p]
  ```

  The KL is "reverse" because in maximum likelihood we think of minimizing `q`
  as in `KL[p, q]`.

  Warning: when self_normalized = True` this function makes non-log-space
  calculations and may therefore be numerically unstable for `|logu| >> 0`.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    self_normalized: Python `bool` indicating whether `f'(u=1)=0`. When
      `f'(u=1)=0` the implied Csiszar f-Divergence remains non-negative even
      when `p, q` are unnormalized measures.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    kl_reverse_of_u: `float`-like `Tensor` of the Csiszar-function evaluated at
      `u = exp(logu)`.

  Raises:
    TypeError: if `self_normalized` is `None` or a `Tensor`.
  """

  with tf.compat.v1.name_scope(name, "kl_reverse", [logu]):
    return amari_alpha(logu, alpha=0., self_normalized=self_normalized)

def kl_reverse ( logu , self_normalized = False , name = None ) : with tf . compat . v1 . name_scope ( name , "kl_reverse" , [ logu ] ) : return amari_alpha ( logu , alpha = 0. , self_normalized = self_normalized )

The reverse Kullback - Leibler Csiszar - function in log - space .


def jensen_shannon(logu, self_normalized=False, name=None):
  """The Jensen-Shannon Csiszar-function in log-space.

  A Csiszar-function is a member of,

  ```none
  F = { f:R_+ to R : f convex }.
  ```

  When `self_normalized = True`, the Jensen-Shannon Csiszar-function is:

  ```none
  f(u) = u log(u) - (1 + u) log(1 + u) + (u + 1) log(2)
  ```

  When `self_normalized = False` the `(u + 1) log(2)` term is omitted.

  Observe that as an f-Divergence, this Csiszar-function implies:

  ```none
  D_f[p, q] = KL[p, m] + KL[q, m]
  m(x) = 0.5 p(x) + 0.5 q(x)
  ```

  In a sense, this divergence is the "reverse" of the Arithmetic-Geometric
  f-Divergence.

  This Csiszar-function induces a symmetric f-Divergence, i.e.,
  `D_f[p, q] = D_f[q, p]`.

  Warning: this function makes non-log-space calculations and may therefore be
  numerically unstable for `|logu| >> 0`.

  For more information, see:
    Lin, J. "Divergence measures based on the Shannon entropy." IEEE Trans.
    Inf. Th., 37, 145-151, 1991.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    self_normalized: Python `bool` indicating whether `f'(u=1)=0`. When
      `f'(u=1)=0` the implied Csiszar f-Divergence remains non-negative even
      when `p, q` are unnormalized measures.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    jensen_shannon_of_u: `float`-like `Tensor` of the Csiszar-function
      evaluated at `u = exp(logu)`.
  """

  with tf.compat.v1.name_scope(name, "jensen_shannon", [logu]):
    logu = tf.convert_to_tensor(value=logu, name="logu")
    npdt = logu.dtype.as_numpy_dtype
    y = tf.nn.softplus(logu)
    if self_normalized:
      y -= np.log(2).astype(npdt)
    return tf.exp(logu) * logu - (1. + tf.exp(logu)) * y

def jensen_shannon ( logu , self_normalized = False , name = None ) : with tf . compat . v1 . name_scope ( name , "jensen_shannon" , [ logu ] ) : logu = tf . convert_to_tensor ( value = logu , name = "logu" ) npdt = logu . dtype . as_numpy_dtype y = tf . nn . softplus ( logu ) if self_normalized : y -= np . log ( 2 ) . astype ( npdt ) return tf . exp ( logu ) * logu - ( 1. + tf . exp ( logu ) ) * y

The Jensen - Shannon Csiszar - function in log - space .


def pearson(logu, name=None):
  """The Pearson Csiszar-function in log-space.

  A Csiszar-function is a member of,

  ```none
  F = { f:R_+ to R : f convex }.
  ```

  The Pearson Csiszar-function is:

  ```none
  f(u) = (u - 1)**2
  ```

  Warning: this function makes non-log-space calculations and may therefore be
  numerically unstable for `|logu| >> 0`.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    pearson_of_u: `float`-like `Tensor` of the Csiszar-function evaluated at
      `u = exp(logu)`.
  """

  with tf.compat.v1.name_scope(name, "pearson", [logu]):
    logu = tf.convert_to_tensor(value=logu, name="logu")
    return tf.square(tf.math.expm1(logu))

def pearson ( logu , name = None ) : with tf . compat . v1 . name_scope ( name , "pearson" , [ logu ] ) : logu = tf . convert_to_tensor ( value = logu , name = "logu" ) return tf . square ( tf . math . expm1 ( logu ) )

The Pearson Csiszar - function in log - space .


def squared_hellinger(logu, name=None):
  """The Squared-Hellinger Csiszar-function in log-space.

  A Csiszar-function is a member of,

  ```none
  F = { f:R_+ to R : f convex }.
  ```

  The Squared-Hellinger Csiszar-function is:

  ```none
  f(u) = (sqrt(u) - 1)**2
  ```

  This Csiszar-function induces a symmetric f-Divergence, i.e.,
  `D_f[p, q] = D_f[q, p]`.

  Warning: this function makes non-log-space calculations and may therefore be
  numerically unstable for `|logu| >> 0`.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    squared_hellinger_of_u: `float`-like `Tensor` of the Csiszar-function
      evaluated at `u = exp(logu)`.
  """

  with tf.compat.v1.name_scope(name, "squared_hellinger", [logu]):
    logu = tf.convert_to_tensor(value=logu, name="logu")
    return pearson(0.5 * logu)

def squared_hellinger ( logu , name = None ) : with tf . compat . v1 . name_scope ( name , "squared_hellinger" , [ logu ] ) : logu = tf . convert_to_tensor ( value = logu , name = "logu" ) return pearson ( 0.5 * logu )

The Squared - Hellinger Csiszar - function in log - space .


def triangular(logu, name=None):
  """The Triangular Csiszar-function in log-space.

  A Csiszar-function is a member of,

  ```none
  F = { f:R_+ to R : f convex }.
  ```

  The Triangular Csiszar-function is:

  ```none
  f(u) = (u - 1)**2 / (1 + u)
  ```

  This Csiszar-function induces a symmetric f-Divergence, i.e.,
  `D_f[p, q] = D_f[q, p]`.

  Warning: this function makes non-log-space calculations and may therefore be
  numerically unstable for `|logu| >> 0`.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    triangular_of_u: `float`-like `Tensor` of the Csiszar-function evaluated
      at `u = exp(logu)`.
  """

  with tf.compat.v1.name_scope(name, "triangular", [logu]):
    logu = tf.convert_to_tensor(value=logu, name="logu")
    return pearson(logu) / (1. + tf.exp(logu))

def triangular ( logu , name = None ) : with tf . compat . v1 . name_scope ( name , "triangular" , [ logu ] ) : logu = tf . convert_to_tensor ( value = logu , name = "logu" ) return pearson ( logu ) / ( 1. + tf . exp ( logu ) )

The Triangular Csiszar - function in log - space .


def t_power(logu, t, self_normalized=False, name=None):
  """The T-Power Csiszar-function in log-space.

  A Csiszar-function is a member of,

  ```none
  F = { f:R_+ to R : f convex }.
  ```

  When `self_normalized = True` the T-Power Csiszar-function is:

  ```none
  f(u) = s [ u**t - 1 - t(u - 1) ]
  s = { -1   0 < t < 1
      { +1   otherwise
  ```

  When `self_normalized = False` the `- t(u - 1)` term is omitted.

  This is similar to the `amari_alpha` Csiszar-function, with the associated
  divergence being the same up to factors depending only on `t`.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    t:  `Tensor` of same `dtype` as `logu` and broadcastable shape.
    self_normalized: Python `bool` indicating whether `f'(u=1)=0`.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    t_power_of_u: `float`-like `Tensor` of the Csiszar-function evaluated
      at `u = exp(logu)`.
  """
  with tf.compat.v1.name_scope(name, "t_power", [logu, t]):
    logu = tf.convert_to_tensor(value=logu, name="logu")
    t = tf.convert_to_tensor(value=t, dtype=logu.dtype.base_dtype, name="t")
    fu = tf.math.expm1(t * logu)
    if self_normalized:
      fu -= t * tf.math.expm1(logu)
    fu *= tf.where(tf.logical_and(0. < t, t < 1.),
                   -tf.ones_like(t),
                   tf.ones_like(t))
    return fu

def t_power ( logu , t , self_normalized = False , name = None ) : with tf . compat . v1 . name_scope ( name , "t_power" , [ logu , t ] ) : logu = tf . convert_to_tensor ( value = logu , name = "logu" ) t = tf . convert_to_tensor ( value = t , dtype = logu . dtype . base_dtype , name = "t" ) fu = tf . math . expm1 ( t * logu ) if self_normalized : fu -= t * tf . math . expm1 ( logu ) fu *= tf . where ( tf . logical_and ( 0. < t , t < 1. ) , - tf . ones_like ( t ) , tf . ones_like ( t ) ) return fu

The T - Power Csiszar - function in log - space .


def log1p_abs(logu, name=None):
  """The log1p-abs Csiszar-function in log-space.

  A Csiszar-function is a member of,

  ```none
  F = { f:R_+ to R : f convex }.
  ```

  The Log1p-Abs Csiszar-function is:

  ```none
  f(u) = u**(sign(u-1)) - 1
  ```

  This function is so-named because it was invented from the following recipe.
  Choose a convex function g such that g(0)=0 and solve for f:

  ```none
  log(1 + f(u)) = g(log(u)).
    <=>
  f(u) = exp(g(log(u))) - 1
  ```

  That is, the graph is identically `g` when y-axis is `log1p`-domain and x-axis
  is `log`-domain.

  Warning: this function makes non-log-space calculations and may therefore be
  numerically unstable for `|logu| >> 0`.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    log1p_abs_of_u: `float`-like `Tensor` of the Csiszar-function evaluated
      at `u = exp(logu)`.
  """

  with tf.compat.v1.name_scope(name, "log1p_abs", [logu]):
    logu = tf.convert_to_tensor(value=logu, name="logu")
    return tf.math.expm1(tf.abs(logu))

def log1p_abs ( logu , name = None ) : with tf . compat . v1 . name_scope ( name , "log1p_abs" , [ logu ] ) : logu = tf . convert_to_tensor ( value = logu , name = "logu" ) return tf . math . expm1 ( tf . abs ( logu ) )

The log1p - abs Csiszar - function in log - space .


def jeffreys(logu, name=None):
  """The Jeffreys Csiszar-function in log-space.

  A Csiszar-function is a member of,

  ```none
  F = { f:R_+ to R : f convex }.
  ```

  The Jeffreys Csiszar-function is:

  ```none
  f(u) = 0.5 ( u log(u) - log(u) )
       = 0.5 kl_forward + 0.5 kl_reverse
       = symmetrized_csiszar_function(kl_reverse)
       = symmetrized_csiszar_function(kl_forward)
  ```

  This Csiszar-function induces a symmetric f-Divergence, i.e.,
  `D_f[p, q] = D_f[q, p]`.

  Warning: this function makes non-log-space calculations and may therefore be
  numerically unstable for `|logu| >> 0`.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    jeffreys_of_u: `float`-like `Tensor` of the Csiszar-function evaluated
      at `u = exp(logu)`.
  """

  with tf.compat.v1.name_scope(name, "jeffreys", [logu]):
    logu = tf.convert_to_tensor(value=logu, name="logu")
    return 0.5 * tf.math.expm1(logu) * logu

def jeffreys ( logu , name = None ) : with tf . compat . v1 . name_scope ( name , "jeffreys" , [ logu ] ) : logu = tf . convert_to_tensor ( value = logu , name = "logu" ) return 0.5 * tf . math . expm1 ( logu ) * logu

The Jeffreys Csiszar - function in log - space .


def modified_gan(logu, self_normalized=False, name=None):
  """The Modified-GAN Csiszar-function in log-space.

  A Csiszar-function is a member of,

  ```none
  F = { f:R_+ to R : f convex }.
  ```

  When `self_normalized = True` the modified-GAN (Generative/Adversarial
  Network) Csiszar-function is:

  ```none
  f(u) = log(1 + u) - log(u) + 0.5 (u - 1)
  ```

  When `self_normalized = False` the `0.5 (u - 1)` is omitted.

  The unmodified GAN Csiszar-function is identical to Jensen-Shannon (with
  `self_normalized = False`).

  Warning: this function makes non-log-space calculations and may therefore be
  numerically unstable for `|logu| >> 0`.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    self_normalized: Python `bool` indicating whether `f'(u=1)=0`. When
      `f'(u=1)=0` the implied Csiszar f-Divergence remains non-negative even
      when `p, q` are unnormalized measures.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    chi_square_of_u: `float`-like `Tensor` of the Csiszar-function evaluated
      at `u = exp(logu)`.
  """

  with tf.compat.v1.name_scope(name, "chi_square", [logu]):
    logu = tf.convert_to_tensor(value=logu, name="logu")
    y = tf.nn.softplus(logu) - logu
    if self_normalized:
      y += 0.5 * tf.math.expm1(logu)
    return y

def modified_gan ( logu , self_normalized = False , name = None ) : with tf . compat . v1 . name_scope ( name , "chi_square" , [ logu ] ) : logu = tf . convert_to_tensor ( value = logu , name = "logu" ) y = tf . nn . softplus ( logu ) - logu if self_normalized : y += 0.5 * tf . math . expm1 ( logu ) return y

The Modified - GAN Csiszar - function in log - space .


def dual_csiszar_function(logu, csiszar_function, name=None):
  """Calculates the dual Csiszar-function in log-space.

  A Csiszar-function is a member of,

  ```none
  F = { f:R_+ to R : f convex }.
  ```

  The Csiszar-dual is defined as:

  ```none
  f^*(u) = u f(1 / u)
  ```

  where `f` is some other Csiszar-function.

  For example, the dual of `kl_reverse` is `kl_forward`, i.e.,

  ```none
  f(u) = -log(u)
  f^*(u) = u f(1 / u) = -u log(1 / u) = u log(u)
  ```

  The dual of the dual is the original function:

  ```none
  f^**(u) = {u f(1/u)}^*(u) = u (1/u) f(1/(1/u)) = f(u)
  ```

  Warning: this function makes non-log-space calculations and may therefore be
  numerically unstable for `|logu| >> 0`.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    csiszar_function: Python `callable` representing a Csiszar-function over
      log-domain.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    dual_f_of_u: `float`-like `Tensor` of the result of calculating the dual of
      `f` at `u = exp(logu)`.
  """

  with tf.compat.v1.name_scope(name, "dual_csiszar_function", [logu]):
    return tf.exp(logu) * csiszar_function(-logu)

def dual_csiszar_function ( logu , csiszar_function , name = None ) : with tf . compat . v1 . name_scope ( name , "dual_csiszar_function" , [ logu ] ) : return tf . exp ( logu ) * csiszar_function ( - logu )

Calculates the dual Csiszar - function in log - space .


def symmetrized_csiszar_function(logu, csiszar_function, name=None):
  """Symmetrizes a Csiszar-function in log-space.

  A Csiszar-function is a member of,

  ```none
  F = { f:R_+ to R : f convex }.
  ```

  The symmetrized Csiszar-function is defined as:

  ```none
  f_g(u) = 0.5 g(u) + 0.5 u g (1 / u)
  ```

  where `g` is some other Csiszar-function.

  We say the function is "symmetrized" because:

  ```none
  D_{f_g}[p, q] = D_{f_g}[q, p]
  ```

  for all `p << >> q` (i.e., `support(p) = support(q)`).

  There exists alternatives for symmetrizing a Csiszar-function. For example,

  ```none
  f_g(u) = max(f(u), f^*(u)),
  ```

  where `f^*` is the dual Csiszar-function, also implies a symmetric
  f-Divergence.

  Example:

  When either of the following functions are symmetrized, we obtain the
  Jensen-Shannon Csiszar-function, i.e.,

  ```none
  g(u) = -log(u) - (1 + u) log((1 + u) / 2) + u - 1
  h(u) = log(4) + 2 u log(u / (1 + u))
  ```

  implies,

  ```none
  f_g(u) = f_h(u) = u log(u) - (1 + u) log((1 + u) / 2)
         = jensen_shannon(log(u)).
  ```

  Warning: this function makes non-log-space calculations and may therefore be
  numerically unstable for `|logu| >> 0`.

  Args:
    logu: `float`-like `Tensor` representing `log(u)` from above.
    csiszar_function: Python `callable` representing a Csiszar-function over
      log-domain.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    symmetrized_g_of_u: `float`-like `Tensor` of the result of applying the
      symmetrization of `g` evaluated at `u = exp(logu)`.
  """

  with tf.compat.v1.name_scope(name, "symmetrized_csiszar_function", [logu]):
    logu = tf.convert_to_tensor(value=logu, name="logu")
    return 0.5 * (csiszar_function(logu)
                  + dual_csiszar_function(logu, csiszar_function))

def symmetrized_csiszar_function ( logu , csiszar_function , name = None ) : with tf . compat . v1 . name_scope ( name , "symmetrized_csiszar_function" , [ logu ] ) : logu = tf . convert_to_tensor ( value = logu , name = "logu" ) return 0.5 * ( csiszar_function ( logu ) + dual_csiszar_function ( logu , csiszar_function ) )

Symmetrizes a Csiszar - function in log - space .


def monte_carlo_csiszar_f_divergence(
    f,
    p_log_prob,
    q,
    num_draws,
    use_reparametrization=None,
    seed=None,
    name=None):
  """Monte-Carlo approximation of the Csiszar f-Divergence.

  A Csiszar-function is a member of,

  ```none
  F = { f:R_+ to R : f convex }.
  ```

  The Csiszar f-Divergence for Csiszar-function f is given by:

  ```none
  D_f[p(X), q(X)] := E_{q(X)}[ f( p(X) / q(X) ) ]
                  ~= m**-1 sum_j^m f( p(x_j) / q(x_j) ),
                             where x_j ~iid q(X)
  ```

  Tricks: Reparameterization and Score-Gradient

  When q is "reparameterized", i.e., a diffeomorphic transformation of a
  parameterless distribution (e.g.,
  `Normal(Y; m, s) <=> Y = sX + m, X ~ Normal(0,1)`), we can swap gradient and
  expectation, i.e.,
  `grad[Avg{ s_i : i=1...n }] = Avg{ grad[s_i] : i=1...n }` where `S_n=Avg{s_i}`
  and `s_i = f(x_i), x_i ~iid q(X)`.

  However, if q is not reparameterized, TensorFlow's gradient will be incorrect
  since the chain-rule stops at samples of unreparameterized distributions. In
  this circumstance using the Score-Gradient trick results in an unbiased
  gradient, i.e.,

  ```none
  grad[ E_q[f(X)] ]
  = grad[ int dx q(x) f(x) ]
  = int dx grad[ q(x) f(x) ]
  = int dx [ q'(x) f(x) + q(x) f'(x) ]
  = int dx q(x) [q'(x) / q(x) f(x) + f'(x) ]
  = int dx q(x) grad[ f(x) q(x) / stop_grad[q(x)] ]
  = E_q[ grad[ f(x) q(x) / stop_grad[q(x)] ] ]
  ```

  Unless `q.reparameterization_type != tfd.FULLY_REPARAMETERIZED` it is
  usually preferable to set `use_reparametrization = True`.

  Example Application:

  The Csiszar f-Divergence is a useful framework for variational inference.
  I.e., observe that,

  ```none
  f(p(x)) =  f( E_{q(Z | x)}[ p(x, Z) / q(Z | x) ] )
          <= E_{q(Z | x)}[ f( p(x, Z) / q(Z | x) ) ]
          := D_f[p(x, Z), q(Z | x)]
  ```

  The inequality follows from the fact that the "perspective" of `f`, i.e.,
  `(s, t) |-> t f(s / t))`, is convex in `(s, t)` when `s/t in domain(f)` and
  `t` is a real. Since the above framework includes the popular Evidence Lower
  BOund (ELBO) as a special case, i.e., `f(u) = -log(u)`, we call this framework
  "Evidence Divergence Bound Optimization" (EDBO).

  Args:
    f: Python `callable` representing a Csiszar-function in log-space, i.e.,
      takes `p_log_prob(q_samples) - q.log_prob(q_samples)`.
    p_log_prob: Python `callable` taking (a batch of) samples from `q` and
      returning the natural-log of the probability under distribution `p`.
      (In variational inference `p` is the joint distribution.)
    q: `tf.Distribution`-like instance; must implement:
      `reparameterization_type`, `sample(n, seed)`, and `log_prob(x)`.
      (In variational inference `q` is the approximate posterior distribution.)
    num_draws: Integer scalar number of draws used to approximate the
      f-Divergence expectation.
    use_reparametrization: Python `bool`. When `None` (the default),
      automatically set to:
      `q.reparameterization_type == tfd.FULLY_REPARAMETERIZED`.
      When `True` uses the standard Monte-Carlo average. When `False` uses the
      score-gradient trick. (See above for details.)  When `False`, consider
      using `csiszar_vimco`.
    seed: Python `int` seed for `q.sample`.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    monte_carlo_csiszar_f_divergence: `float`-like `Tensor` Monte Carlo
      approximation of the Csiszar f-Divergence.

  Raises:
    ValueError: if `q` is not a reparameterized distribution and
      `use_reparametrization = True`. A distribution `q` is said to be
      "reparameterized" when its samples are generated by transforming the
      samples of another distribution which does not depend on the
      parameterization of `q`. This property ensures the gradient (with respect
      to parameters) is valid.
    TypeError: if `p_log_prob` is not a Python `callable`.
  """
  reparameterization_types = tf.nest.flatten(q.reparameterization_type)
  with tf.compat.v1.name_scope(name, "monte_carlo_csiszar_f_divergence",
                               [num_draws]):
    if use_reparametrization is None:
      use_reparametrization = all(
          reparameterization_type == tfd.FULLY_REPARAMETERIZED
          for reparameterization_type in reparameterization_types)
    elif (use_reparametrization and
          any(reparameterization_type != tfd.FULLY_REPARAMETERIZED
              for reparameterization_type in reparameterization_types)):
      # TODO(jvdillon): Consider only raising an exception if the gradient is
      # requested.
      raise ValueError(
          "Distribution `q` must be reparameterized, i.e., a diffeomorphic "
          "transformation of a parameterless distribution. (Otherwise this "
          "function has a biased gradient.)")
    if not callable(p_log_prob):
      raise TypeError("`p_log_prob` must be a Python `callable` function.")
    return monte_carlo.expectation(
        f=lambda q_samples: f(p_log_prob(q_samples) - q.log_prob(q_samples)),
        samples=q.sample(num_draws, seed=seed),
        log_prob=q.log_prob,  # Only used if use_reparametrization=False.
        use_reparametrization=use_reparametrization)

Monte - Carlo approximation of the Csiszar f - Divergence .


def csiszar_vimco_helper(logu, name=None):
  """Helper to `csiszar_vimco`; computes `log_avg_u`, `log_sooavg_u`.

  `axis = 0` of `logu` is presumed to correspond to iid samples from `q`, i.e.,

  ```none
  logu[j] = log(u[j])
  u[j] = p(x, h[j]) / q(h[j] | x)
  h[j] iid~ q(H | x)
  ```

  Args:
    logu: Floating-type `Tensor` representing `log(p(x, h) / q(h | x))`.
    name: Python `str` name prefixed to Ops created by this function.

  Returns:
    log_avg_u: `logu.dtype` `Tensor` corresponding to the natural-log of the
      average of `u`. The sum of the gradient of `log_avg_u` is `1`.
    log_sooavg_u: `logu.dtype` `Tensor` characterized by the natural-log of the
      average of `u`` except that the average swaps-out `u[i]` for the
      leave-`i`-out Geometric-average. The mean of the gradient of
      `log_sooavg_u` is `1`. Mathematically `log_sooavg_u` is,
      ```none
      log_sooavg_u[i] = log(Avg{h[j ; i] : j=0, ..., m-1})
      h[j ; i] = { u[j]                              j!=i
                 { GeometricAverage{u[k] : k != i}   j==i
      ```

  """
  with tf.compat.v1.name_scope(name, "csiszar_vimco_helper", [logu]):
    logu = tf.convert_to_tensor(value=logu, name="logu")

    n = tf.compat.dimension_value(logu.shape.with_rank_at_least(1)[0])
    if n is None:
      n = tf.shape(input=logu)[0]
      log_n = tf.math.log(tf.cast(n, dtype=logu.dtype))
      nm1 = tf.cast(n - 1, dtype=logu.dtype)
    else:
      log_n = np.log(n).astype(logu.dtype.as_numpy_dtype)
      nm1 = np.asarray(n - 1, dtype=logu.dtype.as_numpy_dtype)

    # Throughout we reduce across axis=0 since this is presumed to be iid
    # samples.

    log_max_u = tf.reduce_max(input_tensor=logu, axis=0)
    log_sum_u_minus_log_max_u = tf.reduce_logsumexp(
        input_tensor=logu - log_max_u, axis=0)

    # log_loosum_u[i] =
    # = logsumexp(logu[j] : j != i)
    # = log( exp(logsumexp(logu)) - exp(logu[i]) )
    # = log( exp(logsumexp(logu - logu[i])) exp(logu[i])  - exp(logu[i]))
    # = logu[i] + log(exp(logsumexp(logu - logu[i])) - 1)
    # = logu[i] + log(exp(logsumexp(logu) - logu[i]) - 1)
    # = logu[i] + softplus_inverse(logsumexp(logu) - logu[i])
    d = log_sum_u_minus_log_max_u + (log_max_u - logu)
    # We use `d != 0` rather than `d > 0.` because `d < 0.` should never
    # happens; if it does we want to complain loudly (which `softplus_inverse`
    # will).
    d_ok = tf.not_equal(d, 0.)
    safe_d = tf.where(d_ok, d, tf.ones_like(d))
    d_ok_result = logu + tfd.softplus_inverse(safe_d)

    inf = np.array(np.inf, dtype=logu.dtype.as_numpy_dtype)

    # When not(d_ok) and is_positive_and_largest then we manually compute the
    # log_loosum_u. (We can efficiently do this for any one point but not all,
    # hence we still need the above calculation.) This is good because when
    # this condition is met, we cannot use the above calculation; its -inf.
    # We now compute the log-leave-out-max-sum, replicate it to every
    # point and make sure to select it only when we need to.
    is_positive_and_largest = tf.logical_and(
        logu > 0.,
        tf.equal(logu, log_max_u[tf.newaxis, ...]))
    log_lomsum_u = tf.reduce_logsumexp(
        input_tensor=tf.where(is_positive_and_largest,
                              tf.fill(tf.shape(input=logu), -inf), logu),
        axis=0,
        keepdims=True)
    log_lomsum_u = tf.tile(
        log_lomsum_u,
        multiples=1 + tf.pad(tensor=[n - 1], paddings=[[0, tf.rank(logu) - 1]]))

    d_not_ok_result = tf.where(is_positive_and_largest, log_lomsum_u,
                               tf.fill(tf.shape(input=d), -inf))

    log_loosum_u = tf.where(d_ok, d_ok_result, d_not_ok_result)

    # The swap-one-out-sum ("soosum") is n different sums, each of which
    # replaces the i-th item with the i-th-left-out average, i.e.,
    # soo_sum_u[i] = [exp(logu) - exp(logu[i])] + exp(mean(logu[!=i]))
    #              =  exp(log_loosum_u[i])      + exp(looavg_logu[i])
    looavg_logu = (tf.reduce_sum(input_tensor=logu, axis=0) - logu) / nm1
    log_soosum_u = tf.reduce_logsumexp(
        input_tensor=tf.stack([log_loosum_u, looavg_logu]), axis=0)

    log_avg_u = log_sum_u_minus_log_max_u + log_max_u - log_n
    log_sooavg_u = log_soosum_u - log_n

    log_avg_u.set_shape(logu.shape.with_rank_at_least(1)[1:])
    log_sooavg_u.set_shape(logu.shape)

    return log_avg_u, log_sooavg_u


def _assert_ndims_statically(x,
                             expect_ndims=None,
                             expect_ndims_at_least=None,
                             expect_static=False):
  """Assert that Tensor x has expected number of dimensions."""
  ndims = x.shape.ndims
  if ndims is None:
    if expect_static:
      raise ValueError('Expected static ndims. Found: {}'.format(x))
    return
  if expect_ndims is not None and ndims != expect_ndims:
    raise ValueError('ndims must be {}.  Found: {}'.format(expect_ndims, ndims))
  if expect_ndims_at_least is not None and ndims < expect_ndims_at_least:
    raise ValueError('ndims must be at least {}. Found {}'.format(
        expect_ndims_at_least, ndims))

def
_assert_ndims_statically
(
x
,
expect_ndims
=
,
expect_ndims_at_least
=
,
expect_static
=
False
)
:
ndims
=
x
.
shape
.
ndims
if
ndims
is
:
if
expect_static
:
raise
ValueError
(
'Expected static ndims. Found: {}'
.
format
(
x
)
)
return
if
expect_ndims
is
not
and
ndims
!=
expect_ndims
:
raise
ValueError
(
'ndims must be {}.  Found: {}'
.
format
(
expect_ndims
,
ndims
)
)
if
expect_ndims_at_least
is
not
and
ndims
<
expect_ndims_at_least
:
raise
ValueError
(
'ndims must be at least {}. Found {}'
.
format
(
expect_ndims_at_least
,
ndims
)
)

Assert that Tensor x has expected number of dimensions .


def _batch_gather_with_broadcast(params, indices, axis):
  """Like batch_gather, but broadcasts to the left of axis."""
  # batch_gather assumes...
  #   params.shape =  [A1,...,AN, B1,...,BM]
  #   indices.shape = [A1,...,AN, C]
  # which gives output of shape
  #                   [A1,...,AN, C, B1,...,BM]
  # Here we broadcast dims of each to the left of `axis` in params, and left of
  # the rightmost dim in indices, e.g. we can
  # have
  #   params.shape =  [A1,...,AN, B1,...,BM]
  #   indices.shape = [a1,...,aN, C],
  # where ai broadcasts with Ai.

  # leading_bcast_shape is the broadcast of [A1,...,AN] and [a1,...,aN].
  leading_bcast_shape = tf.broadcast_dynamic_shape(
      tf.shape(input=params)[:axis],
      tf.shape(input=indices)[:-1])
  params += tf.zeros(
      tf.concat((leading_bcast_shape, tf.shape(input=params)[axis:]), axis=0),
      dtype=params.dtype)
  indices += tf.zeros(
      tf.concat((leading_bcast_shape, tf.shape(input=indices)[-1:]), axis=0),
      dtype=indices.dtype)
  return tf.compat.v1.batch_gather(params, indices)

Like batch_gather but broadcasts to the left of axis .


def _broadcast_cat_event_and_params(event, params, base_dtype):
  """Broadcasts the event or distribution parameters."""
  if dtype_util.is_integer(event.dtype):
    pass
  elif dtype_util.is_floating(event.dtype):
    # When `validate_args=True` we've already ensured int/float casting
    # is closed.
    event = tf.cast(event, dtype=tf.int32)
  else:
    raise TypeError("`value` should have integer `dtype` or "
                    "`self.dtype` ({})".format(base_dtype))
  shape_known_statically = (
      tensorshape_util.rank(params.shape) is not None and
      tensorshape_util.is_fully_defined(params.shape[:-1]) and
      tensorshape_util.is_fully_defined(event.shape))
  if not shape_known_statically or params.shape[:-1] != event.shape:
    params *= tf.ones_like(event[..., tf.newaxis],
                           dtype=params.dtype)
    params_shape = tf.shape(input=params)[:-1]
    event *= tf.ones(params_shape, dtype=event.dtype)
    if tensorshape_util.rank(params.shape) is not None:
      tensorshape_util.set_shape(event, params.shape[:-1])

  return event, params

Broadcasts the event or distribution parameters .


def expectation_importance_sampler_logspace(
    log_f,
    log_p,
    sampling_dist_q,
    z=None,
    n=None,
    seed=None,
    name='expectation_importance_sampler_logspace'):
  r"""Importance sampling with a positive function, in log-space.

  With \\(p(z) := exp^{log_p(z)}\\), and \\(f(z) = exp{log_f(z)}\\),
  this `Op` returns

  \\(Log[ n^{-1} sum_{i=1}^n [ f(z_i) p(z_i) / q(z_i) ] ],  z_i ~ q,\\)
  \\(\approx Log[ E_q[ f(Z) p(Z) / q(Z) ] ]\\)
  \\(=       Log[E_p[f(Z)]]\\)

  This integral is done in log-space with max-subtraction to better handle the
  often extreme values that `f(z) p(z) / q(z)` can take on.

  In contrast to `expectation_importance_sampler`, this `Op` returns values in
  log-space.

  User supplies either `Tensor` of samples `z`, or number of samples to draw `n`

  Args:
    log_f: Callable mapping samples from `sampling_dist_q` to `Tensors` with
      shape broadcastable to `q.batch_shape`.
      For example, `log_f` works "just like" `sampling_dist_q.log_prob`.
    log_p:  Callable mapping samples from `sampling_dist_q` to `Tensors` with
      shape broadcastable to `q.batch_shape`.
      For example, `log_p` works "just like" `q.log_prob`.
    sampling_dist_q:  The sampling distribution.
      `tfp.distributions.Distribution`.
      `float64` `dtype` recommended.
      `log_p` and `q` should be supported on the same set.
    z:  `Tensor` of samples from `q`, produced by `q.sample` for some `n`.
    n:  Integer `Tensor`.  Number of samples to generate if `z` is not provided.
    seed:  Python integer to seed the random number generator.
    name:  A name to give this `Op`.

  Returns:
    Logarithm of the importance sampling estimate.  `Tensor` with `shape` equal
      to batch shape of `q`, and `dtype` = `q.dtype`.
  """
  q = sampling_dist_q
  with tf.name_scope(name):
    z = _get_samples(q, z, n, seed)
    log_values = log_f(z) + log_p(z) - q.log_prob(z)
    return _logspace_mean(log_values)

r Importance sampling with a positive function in log - space .


def _broadcast_event_and_samples(event, samples, event_ndims):
  """Broadcasts the event or samples."""
  # This is the shape of self.samples, without the samples axis, i.e. the shape
  # of the result of a call to dist.sample(). This way we can broadcast it with
  # event to get a properly-sized event, then add the singleton dim back at
  # -event_ndims - 1.
  samples_shape = tf.concat(
      [tf.shape(input=samples)[:-event_ndims - 1],
       tf.shape(input=samples)[tf.rank(samples) - event_ndims:]],
      axis=0)
  event *= tf.ones(samples_shape, dtype=event.dtype)
  event = tf.expand_dims(event, axis=-event_ndims - 1)
  samples *= tf.ones_like(event, dtype=samples.dtype)

  return event, samples


def minimize(value_and_gradients_function,
             initial_position,
             tolerance=1e-8,
             x_tolerance=0,
             f_relative_tolerance=0,
             initial_inverse_hessian_estimate=None,
             max_iterations=50,
             parallel_iterations=1,
             stopping_condition=None,
             name=None):
  """Applies the BFGS algorithm to minimize a differentiable function.

  Performs unconstrained minimization of a differentiable function using the
  BFGS scheme. For details of the algorithm, see [Nocedal and Wright(2006)][1].

  ### Usage:

  The following example demonstrates the BFGS optimizer attempting to find the
  minimum for a simple two dimensional quadratic objective function.

  ```python
    minimum = np.array([1.0, 1.0])  # The center of the quadratic bowl.
    scales = np.array([2.0, 3.0])  # The scales along the two axes.

    # The objective function and the gradient.
    def quadratic(x):
      value = tf.reduce_sum(scales * (x - minimum) ** 2)
      return value, tf.gradients(value, x)[0]

    start = tf.constant([0.6, 0.8])  # Starting point for the search.
    optim_results = tfp.optimizer.bfgs_minimize(
        quadratic, initial_position=start, tolerance=1e-8)

    with tf.Session() as session:
      results = session.run(optim_results)
      # Check that the search converged
      assert(results.converged)
      # Check that the argmin is close to the actual value.
      np.testing.assert_allclose(results.position, minimum)
      # Print out the total number of function evaluations it took. Should be 6.
      print ("Function evaluations: %d" % results.num_objective_evaluations)
  ```

  ### References:
  [1]: Jorge Nocedal, Stephen Wright. Numerical Optimization. Springer Series in
    Operations Research. pp 136-140. 2006
    http://pages.mtu.edu/~struther/Courses/OLD/Sp2013/5630/Jorge_Nocedal_Numerical_optimization_267490.pdf

  Args:
    value_and_gradients_function:  A Python callable that accepts a point as a
      real `Tensor` and returns a tuple of `Tensor`s of real dtype containing
      the value of the function and its gradient at that point. The function
      to be minimized. The input should be of shape `[..., n]`, where `n` is
      the size of the domain of input points, and all others are batching
      dimensions. The first component of the return value should be a real
      `Tensor` of matching shape `[...]`. The second component (the gradient)
      should also be of shape `[..., n]` like the input value to the function.
    initial_position: real `Tensor` of shape `[..., n]`. The starting point, or
      points when using batching dimensions, of the search procedure. At these
      points the function value and the gradient norm should be finite.
    tolerance: Scalar `Tensor` of real dtype. Specifies the gradient tolerance
      for the procedure. If the supremum norm of the gradient vector is below
      this number, the algorithm is stopped.
    x_tolerance: Scalar `Tensor` of real dtype. If the absolute change in the
      position between one iteration and the next is smaller than this number,
      the algorithm is stopped.
    f_relative_tolerance: Scalar `Tensor` of real dtype. If the relative change
      in the objective value between one iteration and the next is smaller
      than this value, the algorithm is stopped.
    initial_inverse_hessian_estimate: Optional `Tensor` of the same dtype
      as the components of the output of the `value_and_gradients_function`.
      If specified, the shape should broadcastable to shape `[..., n, n]`; e.g.
      if a single `[n, n]` matrix is provided, it will be automatically
      broadcasted to all batches. Alternatively, one can also specify a
      different hessian estimate for each batch member.
      For the correctness of the algorithm, it is required that this parameter
      be symmetric and positive definite. Specifies the starting estimate for
      the inverse of the Hessian at the initial point. If not specified,
      the identity matrix is used as the starting estimate for the
      inverse Hessian.
    max_iterations: Scalar positive int32 `Tensor`. The maximum number of
      iterations for BFGS updates.
    parallel_iterations: Positive integer. The number of iterations allowed to
      run in parallel.
    stopping_condition: (Optional) 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. The default is tfp.optimizer.converged_all
      which only stops when all batch members have either converged or failed.
      An alternative is tfp.optimizer.converged_any which stops as soon as one
      batch member has converged, or when all have failed.
    name: (Optional) Python str. The name prefixed to the ops created by this
      function. If not supplied, the default name 'minimize' is used.

  Returns:
    optimizer_results: A namedtuple containing the following items:
      converged: boolean tensor of shape `[...]` indicating for each batch
        member whether the minimum was found within tolerance.
      failed:  boolean tensor of shape `[...]` indicating for each batch
        member whether a line search step failed to find a suitable step size
        satisfying Wolfe conditions. In the absence of any constraints on the
        number of objective evaluations permitted, this value will
        be the complement of `converged`. However, if there is
        a constraint and the search stopped due to available
        evaluations being exhausted, both `failed` and `converged`
        will be simultaneously False.
      num_objective_evaluations: The total number of objective
        evaluations performed.
      position: A tensor of shape `[..., n]` containing the last argument value
        found during the search from each starting point. If the search
        converged, then this value is the argmin of the objective function.
      objective_value: A tensor of shape `[...]` with the value of the
        objective function at the `position`. If the search converged, then
        this is the (local) minimum of the objective function.
      objective_gradient: A tensor of shape `[..., n]` containing the gradient
        of the objective function at the `position`. If the search converged
        the max-norm of this tensor should be below the tolerance.
      inverse_hessian_estimate: A tensor of shape `[..., n, n]` containing the
        inverse of the estimated Hessian.
  """
  with tf.compat.v1.name_scope(
      name, 'minimize',
      [initial_position, tolerance, initial_inverse_hessian_estimate]):
    initial_position = tf.convert_to_tensor(
        value=initial_position, name='initial_position')
    dtype = initial_position.dtype.base_dtype
    tolerance = tf.convert_to_tensor(
        value=tolerance, dtype=dtype, name='grad_tolerance')
    f_relative_tolerance = tf.convert_to_tensor(
        value=f_relative_tolerance, dtype=dtype, name='f_relative_tolerance')
    x_tolerance = tf.convert_to_tensor(
        value=x_tolerance, dtype=dtype, name='x_tolerance')
    max_iterations = tf.convert_to_tensor(
        value=max_iterations, name='max_iterations')

    input_shape = distribution_util.prefer_static_shape(initial_position)
    batch_shape, domain_size = input_shape[:-1], input_shape[-1]

    if stopping_condition is None:
      stopping_condition = bfgs_utils.converged_all

    # Control inputs are an optional list of tensors to evaluate before
    # the start of the search procedure. These can be used to assert the
    # validity of inputs to the search procedure.
    control_inputs = None

    if initial_inverse_hessian_estimate is None:
      # Create a default initial inverse Hessian.
      initial_inv_hessian = tf.eye(domain_size,
                                   batch_shape=batch_shape,
                                   dtype=dtype,
                                   name='initial_inv_hessian')
    else:
      # If an initial inverse Hessian is supplied, compute some control inputs
      # to ensure that it is positive definite and symmetric.
      initial_inv_hessian = tf.convert_to_tensor(
          value=initial_inverse_hessian_estimate,
          dtype=dtype,
          name='initial_inv_hessian')
      control_inputs = _inv_hessian_control_inputs(initial_inv_hessian)
      hessian_shape = tf.concat([batch_shape, [domain_size, domain_size]], 0)
      initial_inv_hessian = tf.broadcast_to(initial_inv_hessian, hessian_shape)

    # The `state` here is a `BfgsOptimizerResults` tuple with values for the
    # current state of the algorithm computation.
    def _cond(state):
      """Continue if iterations remain and stopping condition is not met."""
      return ((state.num_iterations < max_iterations) &
              tf.logical_not(stopping_condition(state.converged, state.failed)))

    def _body(state):
      """Main optimization loop."""
      search_direction = _get_search_direction(state.inverse_hessian_estimate,
                                               state.objective_gradient)
      derivative_at_start_pt = tf.reduce_sum(
          input_tensor=state.objective_gradient * search_direction, axis=-1)

      # If the derivative at the start point is not negative, recompute the
      # search direction with the initial inverse Hessian.
      needs_reset = (~state.failed & ~state.converged &
                     (derivative_at_start_pt >= 0))

      search_direction_reset = _get_search_direction(
          initial_inv_hessian, state.objective_gradient)

      actual_serch_direction = tf.where(
          needs_reset, search_direction_reset, search_direction)
      actual_inv_hessian = tf.where(
          needs_reset, initial_inv_hessian, state.inverse_hessian_estimate)

      # Replace the hessian estimate in the state, in case it had to be reset.
      current_state = bfgs_utils.update_fields(
          state, inverse_hessian_estimate=actual_inv_hessian)

      next_state = bfgs_utils.line_search_step(
          current_state,
          value_and_gradients_function, actual_serch_direction,
          tolerance, f_relative_tolerance, x_tolerance, stopping_condition)

      # Update the inverse Hessian if needed and continue.
      return [_update_inv_hessian(current_state, next_state)]

    kwargs = bfgs_utils.get_initial_state_args(
        value_and_gradients_function,
        initial_position,
        tolerance,
        control_inputs)
    kwargs['inverse_hessian_estimate'] = initial_inv_hessian
    initial_state = BfgsOptimizerResults(**kwargs)
    return tf.while_loop(
        cond=_cond,
        body=_body,
        loop_vars=[initial_state],
        parallel_iterations=parallel_iterations)[0]

Applies the BFGS algorithm to minimize a differentiable function .


def _inv_hessian_control_inputs(inv_hessian):
  """Computes control inputs to validate a provided inverse Hessian.

  These ensure that the provided inverse Hessian is positive definite and
  symmetric.

  Args:
    inv_hessian: The starting estimate for the inverse of the Hessian at the
      initial point.

  Returns:
    A list of tf.Assert ops suitable for use with tf.control_dependencies.
  """
  # The easiest way to validate if the inverse Hessian is positive definite is
  # to compute its Cholesky decomposition.
  is_positive_definite = tf.reduce_all(
      input_tensor=tf.math.is_finite(tf.linalg.cholesky(inv_hessian)),
      axis=[-1, -2])

  # Then check that the supplied inverse Hessian is symmetric.
  is_symmetric = tf.equal(bfgs_utils.norm(
      inv_hessian - _batch_transpose(inv_hessian), dims=2), 0)

  # Simply adding a control dependencies on these results is not enough to
  # trigger them, we need to add asserts on the results.
  return [tf.Assert(is_positive_definite,
                    ['Initial inverse Hessian is not positive definite.',
                     inv_hessian]),
          tf.Assert(is_symmetric,
                    ['Initial inverse Hessian is not symmetric',
                     inv_hessian])]

Computes control inputs to validate a provided inverse Hessian .


def _update_inv_hessian(prev_state, next_state):
  """Update the BGFS state by computing the next inverse hessian estimate."""
  # Only update the inverse Hessian if not already failed or converged.
  should_update = ~next_state.converged & ~next_state.failed

  # Compute the normalization term (y^T . s), should not update if is singular.
  gradient_delta = next_state.objective_gradient - prev_state.objective_gradient
  position_delta = next_state.position - prev_state.position
  normalization_factor = tf.reduce_sum(
      input_tensor=gradient_delta * position_delta, axis=-1)
  should_update = should_update & ~tf.equal(normalization_factor, 0)

  def _do_update_inv_hessian():
    next_inv_hessian = _bfgs_inv_hessian_update(
        gradient_delta, position_delta, normalization_factor,
        prev_state.inverse_hessian_estimate)
    return bfgs_utils.update_fields(
        next_state,
        inverse_hessian_estimate=tf.where(should_update,
                                          next_inv_hessian,
                                          prev_state.inverse_hessian_estimate))

  return prefer_static.cond(
      tf.reduce_any(input_tensor=should_update),
      _do_update_inv_hessian,
      lambda: next_state)

Update the BGFS state by computing the next inverse hessian estimate .


def _bfgs_inv_hessian_update(grad_delta, position_delta, normalization_factor,
                             inv_hessian_estimate):
  """Applies the BFGS update to the inverse Hessian estimate.

  The BFGS update rule is (note A^T denotes the transpose of a vector/matrix A).

  ```None
    rho = 1/(grad_delta^T * position_delta)
    U = (I - rho * position_delta * grad_delta^T)
    H_1 =  U * H_0 * U^T + rho * position_delta * position_delta^T
  ```

  Here, `H_0` is the inverse Hessian estimate at the previous iteration and
  `H_1` is the next estimate. Note that `*` should be interpreted as the
  matrix multiplication (with the understanding that matrix multiplication for
  scalars is usual multiplication and for matrix with vector is the action of
  the matrix on the vector.).

  The implementation below utilizes an expanded version of the above formula
  to avoid the matrix multiplications that would be needed otherwise. By
  expansion it is easy to see that one only needs matrix-vector or
  vector-vector operations. The expanded version is:

  ```None
    f = 1 + rho * (grad_delta^T * H_0 * grad_delta)
    H_1 - H_0 = - rho * [position_delta * (H_0 * grad_delta)^T +
                        (H_0 * grad_delta) * position_delta^T] +
                  rho * f * [position_delta * position_delta^T]
  ```

  All the terms in square brackets are matrices and are constructed using
  vector outer products. All the other terms on the right hand side are scalars.
  Also worth noting that the first and second lines are both rank 1 updates
  applied to the current inverse Hessian estimate.

  Args:
    grad_delta: Real `Tensor` of shape `[..., n]`. The difference between the
      gradient at the new position and the old position.
    position_delta: Real `Tensor` of shape `[..., n]`. The change in position
      from the previous iteration to the current one.
    normalization_factor: Real `Tensor` of shape `[...]`. Should be equal to
      `grad_delta^T * position_delta`, i.e. `1/rho` as defined above.
    inv_hessian_estimate: Real `Tensor` of shape `[..., n, n]`. The previous
      estimate of the inverse Hessian. Should be positive definite and
      symmetric.

  Returns:
    A tuple containing the following fields
      is_valid: A Boolean `Tensor` of shape `[...]` indicating batch members
        where the update succeeded. The update can fail if the position change
        becomes orthogonal to the gradient change.
      next_inv_hessian_estimate: A `Tensor` of shape `[..., n, n]`. The next
        Hessian estimate updated using the BFGS update scheme. If the
        `inv_hessian_estimate` is symmetric and positive definite, the
        `next_inv_hessian_estimate` is guaranteed to satisfy the same
        conditions.
  """
  # The quadratic form: y^T.H.y; where H is the inverse Hessian and y is the
  # gradient change.
  conditioned_grad_delta = _mul_right(inv_hessian_estimate, grad_delta)
  conditioned_grad_delta_norm = tf.reduce_sum(
      input_tensor=conditioned_grad_delta * grad_delta, axis=-1)

  # The first rank 1 update term requires the outer product: s.y^T.
  cross_term = _tensor_product(position_delta, conditioned_grad_delta)

  def _expand_scalar(s):
    # Expand dimensions of a batch of scalars to multiply or divide a matrix.
    return s[..., tf.newaxis, tf.newaxis]

  # Symmetrize
  cross_term += _tensor_product(conditioned_grad_delta, position_delta)
  position_term = _tensor_product(position_delta, position_delta)
  with tf.control_dependencies([position_term]):
    position_term *= _expand_scalar(
        1 + conditioned_grad_delta_norm / normalization_factor)

  return (inv_hessian_estimate +
          (position_term - cross_term) / _expand_scalar(normalization_factor))

Applies the BFGS update to the inverse Hessian estimate .


def _mul_right(mat, vec):
  """Computes the product of a matrix with a vector on the right.

  Note this supports dynamic shapes and batched computation.

  Examples:

    M = tf.reshape(tf.range(6), shape=(3, 2))
    # => [[0, 1],
    #     [2, 3],
    #     [4, 5]]
    v = tf.constant([1, 2])  # Shape: (2,)
    _mul_right(M, v)
    # => [ 2,  8, 14]  # Shape: (3,)

    M = tf.reshape(tf.range(30), shape=(2, 3, 5))
    # => [[[ 0,  1,  2,  3,  4],
    #     [ 5,  6,  7,  8,  9],
    #     [10, 11, 12, 13, 14]],
    #
    #    [[15, 16, 17, 18, 19],
    #     [20, 21, 22, 23, 24],
    #     [25, 26, 27, 28, 29]]]
    v = tf.reshape(tf.range(10), shape=(2, 5))
    # => [[0, 1, 2, 3, 4],
    #     [5, 6, 7, 8, 9]]
    _mul_right(M, v)
    # => [[ 30,  80, 130],
    #     [605, 780, 955]]  # Shape: (2, 3)

  Args:
    mat: A `tf.Tensor` of shape `[..., n, m]`.
    vec: A `tf.Tensor` of shape `[..., m]`.

  Returns:
    A tensor of shape `[..., n]` with matching batch dimensions.
  """
  return tf.squeeze(tf.matmul(mat, tf.expand_dims(vec, axis=-1)), axis=-1)

def _mul_right ( mat , vec ) : return tf . squeeze ( tf . matmul ( mat , tf . expand_dims ( vec , axis = - 1 ) ) , axis = - 1 )

Computes the product of a matrix with a vector on the right .


def _tensor_product(t1, t2):
  """Computes the outer product of two possibly batched vectors.

  Args:
    t1: A `tf.Tensor` of shape `[..., n]`.
    t2: A `tf.Tensor` of shape `[..., m]`.

  Returns:
    A tensor of shape `[..., n, m]` with matching batch dimensions, let's call
    it `r`, whose components are:

    ```None
      r[..., i, j] = t1[..., i] * t2[..., j]
    ```
  """
  return tf.matmul(tf.expand_dims(t1, axis=-1), tf.expand_dims(t2, axis=-2))

def _tensor_product ( t1 , t2 ) : return tf . matmul ( tf . expand_dims ( t1 , axis = - 1 ) , tf . expand_dims ( t2 , axis = - 2 ) )

Computes the outer product of two possibly batched vectors .


def _batch_transpose(mat):
  """Transpose a possibly batched matrix.

  Args:
    mat: A `tf.Tensor` of shape `[..., n, m]`.

  Returns:
    A tensor of shape `[..., m, n]` with matching batch dimensions.
  """
  n = distribution_util.prefer_static_rank(mat)
  perm = tf.range(n)
  perm = tf.concat([perm[:-2], [perm[-1], perm[-2]]], axis=0)
  return tf.transpose(a=mat, perm=perm)

def _batch_transpose ( mat ) : n = distribution_util . prefer_static_rank ( mat ) perm = tf . range ( n ) perm = tf . concat ( [ perm [ : - 2 ] , [ perm [ - 1 ] , perm [ - 2 ] ] ] , axis = 0 ) return tf . transpose ( a = mat , perm = perm )


def pad_shape_right_with_ones(x, ndims):
  """Maybe add `ndims` ones to `x.shape` on the right.

  If `ndims` is zero, this is a no-op; otherwise, we will create and return a
  new `Tensor` whose shape is that of `x` with `ndims` ones concatenated on the
  right side. If the shape of `x` is known statically, the shape of the return
  value will be as well.

  Args:
    x: The `Tensor` we'll return a reshaping of.
    ndims: Python `integer` number of ones to pad onto `x.shape`.
  Returns:
    If `ndims` is zero, `x`; otherwise, a `Tensor` whose shape is that of `x`
    with `ndims` ones concatenated on the right side. If possible, returns a
    `Tensor` whose shape is known statically.
  Raises:
    ValueError: if `ndims` is not a Python `integer` greater than or equal to
    zero.
  """
  if not (isinstance(ndims, int) and ndims >= 0):
    raise ValueError(
        '`ndims` must be a Python `integer` greater than zero. Got: {}'
        .format(ndims))
  if ndims == 0:
    return x
  x = tf.convert_to_tensor(value=x)
  original_shape = x.shape
  new_shape = distribution_util.pad(
      tf.shape(input=x), axis=0, back=True, value=1, count=ndims)
  x = tf.reshape(x, new_shape)
  x.set_shape(original_shape.concatenate([1]*ndims))
  return x

Maybe add ndims ones to x . shape on the right .


def sum_rightmost_ndims_preserving_shape(x, ndims):
  """Return `Tensor` with right-most ndims summed.

  Args:
    x: the `Tensor` whose right-most `ndims` dimensions to sum
    ndims: number of right-most dimensions to sum.

  Returns:
    A `Tensor` resulting from calling `reduce_sum` on the `ndims` right-most
    dimensions. If the shape of `x` is statically known, the result will also
    have statically known shape. Otherwise, the resulting shape will only be
    known at runtime.
  """
  x = tf.convert_to_tensor(value=x)
  if x.shape.ndims is not None:
    axes = tf.range(x.shape.ndims - ndims, x.shape.ndims)
  else:
    axes = tf.range(tf.rank(x) - ndims, tf.rank(x))
  return tf.reduce_sum(input_tensor=x, axis=axes)

def sum_rightmost_ndims_preserving_shape ( x , ndims ) : x = tf . convert_to_tensor ( value = x ) if x . shape . ndims is not None : axes = tf . range ( x . shape . ndims - ndims , x . shape . ndims ) else : axes = tf . range ( tf . rank ( x ) - ndims , tf . rank ( x ) ) return tf . reduce_sum ( input_tensor = x , axis = axes )

Return Tensor with right - most ndims summed .


def sqrt_with_finite_grads(x, name=None):
  """A sqrt function whose gradient at zero is very large but finite.

  Args:
    x: a `Tensor` whose sqrt is to be computed.
    name: a Python `str` prefixed to all ops created by this function.
      Default `None` (i.e., "sqrt_with_finite_grads").

  Returns:
    sqrt: the square root of `x`, with an overridden gradient at zero
    grad: a gradient function, which is the same as sqrt's gradient everywhere
      except at zero, where it is given a large finite value, instead of `inf`.

  Raises:
    TypeError: if `tf.convert_to_tensor(x)` is not a `float` type.

  Often in kernel functions, we need to compute the L2 norm of the difference
  between two vectors, `x` and `y`: `sqrt(sum_i((x_i - y_i) ** 2))`. In the
  case where `x` and `y` are identical, e.g., on the diagonal of a kernel
  matrix, we get `NaN`s when we take gradients with respect to the inputs. To
  see, this consider the forward pass:

    ```
    [x_1 ... x_N]  -->  [x_1 ** 2 ... x_N ** 2]  -->
        (x_1 ** 2 + ... + x_N ** 2)  -->  sqrt((x_1 ** 2 + ... + x_N ** 2))
    ```

  When we backprop through this forward pass, the `sqrt` yields an `inf` because
  `grad_z(sqrt(z)) = 1 / (2 * sqrt(z))`. Continuing the backprop to the left, at
  the `x ** 2` term, we pick up a `2 * x`, and when `x` is zero, we get
  `0 * inf`, which is `NaN`.

  We'd like to avoid these `NaN`s, since they infect the rest of the connected
  computation graph. Practically, when two inputs to a kernel function are
  equal, we are in one of two scenarios:
    1. We are actually computing k(x, x), in which case norm(x - x) is
       identically zero, independent of x. In this case, we'd like the
       gradient to reflect this independence: it should be zero.
    2. We are computing k(x, y), and x just *happens* to have the same value
       as y. The gradient at such inputs is in fact ill-defined (there is a
       cusp in the sqrt((x - y) ** 2) surface along the line x = y). There are,
       however, an infinite number of sub-gradients, all of which are valid at
       all such inputs. By symmetry, there is exactly one which is "special":
       zero, and we elect to use that value here. In practice, having two
       identical inputs to a kernel matrix is probably a pathological
       situation to be avoided, but that is better resolved at a higher level
       than this.

  To avoid the infinite gradient at zero, we use tf.custom_gradient to redefine
  the gradient at zero. We assign it to be a very large value, specifically
  the sqrt of the max value of the floating point dtype of the input. We use
  the sqrt (as opposed to just using the max floating point value) to avoid
  potential overflow when combining this value with others downstream.
  """
  with tf.compat.v1.name_scope(name, 'sqrt_with_finite_grads', [x]):
    x = tf.convert_to_tensor(value=x, name='x')
    if not x.dtype.is_floating:
      raise TypeError('Input `x` must be floating type.')
    def grad(grad_ys):
      large_float_like_x = np.sqrt(np.finfo(x.dtype.as_numpy_dtype()).max)
      safe_grads = tf.where(
          tf.equal(x, 0), tf.fill(tf.shape(input=x), large_float_like_x),
          0.5 * tf.math.rsqrt(x))
      return grad_ys * safe_grads
    return tf.sqrt(x), grad

def
sqrt_with_finite_grads
(
x
,
name
=
)
:
with
tf
.
compat
.
v1
.
name_scope
(
name
,
'sqrt_with_finite_grads'
,
[
x
]
)
:
x
=
tf
.
convert_to_tensor
(
value
=
x
,
name
=
'x'
)
if
not
x
.
dtype
.
is_floating
:
raise
TypeError
(
'Input `x` must be floating type.'
)
def
grad
(
grad_ys
)
:
large_float_like_x
=
np
.
sqrt
(
np
.
finfo
(
x
.
dtype
.
as_numpy_dtype
(
)
)
.
max
)
safe_grads
=
tf
.
where
(
tf
.
equal
(
x
,
0
)
,
tf
.
fill
(
tf
.
shape
(
input
=
x
)
,
large_float_like_x
)
,
0.5
*
tf
.
math
.
rsqrt
(
x
)
)
return
grad_ys
*
safe_grads
return
tf
.
sqrt
(
x
)
,
grad

A sqrt function whose gradient at zero is very large but finite .


def maybe_get_common_dtype(arg_list):
  """Return common dtype of arg_list, or None.

  Args:
    arg_list: an iterable of items which are either `None` or have a `dtype`
      property.

  Returns:
    dtype: The common dtype of items in `arg_list`, or `None` if the list is
      empty or all items are `None`.
  """
  # Note that `all` defaults to `True` if `arg_list` is empty.
  if all(a is None for a in arg_list):
    return None
  return dtype_util.common_dtype(arg_list, tf.float32)

Return common dtype of arg_list or None .


def minimize(value_and_gradients_function,
             initial_position,
             num_correction_pairs=10,
             tolerance=1e-8,
             x_tolerance=0,
             f_relative_tolerance=0,
             initial_inverse_hessian_estimate=None,
             max_iterations=50,
             parallel_iterations=1,
             stopping_condition=None,
             name=None):
  """Applies the L-BFGS algorithm to minimize a differentiable function.

  Performs unconstrained minimization of a differentiable function using the
  L-BFGS scheme. See [Nocedal and Wright(2006)][1] for details of the algorithm.

  ### Usage:

  The following example demonstrates the L-BFGS optimizer attempting to find the
  minimum for a simple high-dimensional quadratic objective function.

  ```python
    # A high-dimensional quadratic bowl.
    ndims = 60
    minimum = np.ones([ndims], dtype='float64')
    scales = np.arange(ndims, dtype='float64') + 1.0

    # The objective function and the gradient.
    def quadratic(x):
      value = tf.reduce_sum(scales * (x - minimum) ** 2)
      return value, tf.gradients(value, x)[0]

    start = np.arange(ndims, 0, -1, dtype='float64')
    optim_results = tfp.optimizer.lbfgs_minimize(
        quadratic, initial_position=start, num_correction_pairs=10,
        tolerance=1e-8)

    with tf.Session() as session:
      results = session.run(optim_results)
      # Check that the search converged
      assert(results.converged)
      # Check that the argmin is close to the actual value.
      np.testing.assert_allclose(results.position, minimum)
  ```

  ### References:

  [1] Jorge Nocedal, Stephen Wright. Numerical Optimization. Springer Series
      in Operations Research. pp 176-180. 2006

  http://pages.mtu.edu/~struther/Courses/OLD/Sp2013/5630/Jorge_Nocedal_Numerical_optimization_267490.pdf

  Args:
    value_and_gradients_function:  A Python callable that accepts a point as a
      real `Tensor` and returns a tuple of `Tensor`s of real dtype containing
      the value of the function and its gradient at that point. The function
      to be minimized. The input is of shape `[..., n]`, where `n` is the size
      of the domain of input points, and all others are batching dimensions.
      The first component of the return value is a real `Tensor` of matching
      shape `[...]`. The second component (the gradient) is also of shape
      `[..., n]` like the input value to the function.
    initial_position: Real `Tensor` of shape `[..., n]`. The starting point, or
      points when using batching dimensions, of the search procedure. At these
      points the function value and the gradient norm should be finite.
    num_correction_pairs: Positive integer. Specifies the maximum number of
      (position_delta, gradient_delta) correction pairs to keep as implicit
      approximation of the Hessian matrix.
    tolerance: Scalar `Tensor` of real dtype. Specifies the gradient tolerance
      for the procedure. If the supremum norm of the gradient vector is below
      this number, the algorithm is stopped.
    x_tolerance: Scalar `Tensor` of real dtype. If the absolute change in the
      position between one iteration and the next is smaller than this number,
      the algorithm is stopped.
    f_relative_tolerance: Scalar `Tensor` of real dtype. If the relative change
      in the objective value between one iteration and the next is smaller
      than this value, the algorithm is stopped.
    initial_inverse_hessian_estimate: None. Option currently not supported.
    max_iterations: Scalar positive int32 `Tensor`. The maximum number of
      iterations for L-BFGS updates.
    parallel_iterations: Positive integer. The number of iterations allowed to
      run in parallel.
    stopping_condition: (Optional) 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. The default is tfp.optimizer.converged_all
      which only stops when all batch members have either converged or failed.
      An alternative is tfp.optimizer.converged_any which stops as soon as one
      batch member has converged, or when all have failed.
    name: (Optional) Python str. The name prefixed to the ops created by this
      function. If not supplied, the default name 'minimize' is used.

  Returns:
    optimizer_results: A namedtuple containing the following items:
      converged: Scalar boolean tensor indicating whether the minimum was
        found within tolerance.
      failed:  Scalar boolean tensor indicating whether a line search
        step failed to find a suitable step size satisfying Wolfe
        conditions. In the absence of any constraints on the
        number of objective evaluations permitted, this value will
        be the complement of `converged`. However, if there is
        a constraint and the search stopped due to available
        evaluations being exhausted, both `failed` and `converged`
        will be simultaneously False.
      num_objective_evaluations: The total number of objective
        evaluations performed.
      position: A tensor containing the last argument value found
        during the search. If the search converged, then
        this value is the argmin of the objective function.
      objective_value: A tensor containing the value of the objective
        function at the `position`. If the search converged, then this is
        the (local) minimum of the objective function.
      objective_gradient: A tensor containing the gradient of the objective
        function at the `position`. If the search converged the
        max-norm of this tensor should be below the tolerance.
      position_deltas: A tensor encoding information about the latest
        changes in `position` during the algorithm execution.
      gradient_deltas: A tensor encoding information about the latest
        changes in `objective_gradient` during the algorithm execution.
  """
  if initial_inverse_hessian_estimate is not None:
    raise NotImplementedError(
        'Support of initial_inverse_hessian_estimate arg not yet implemented')

  if stopping_condition is None:
    stopping_condition = bfgs_utils.converged_all

  with tf.compat.v1.name_scope(name, 'minimize', [initial_position, tolerance]):
    initial_position = tf.convert_to_tensor(
        value=initial_position, name='initial_position')
    dtype = initial_position.dtype.base_dtype
    tolerance = tf.convert_to_tensor(
        value=tolerance, dtype=dtype, name='grad_tolerance')
    f_relative_tolerance = tf.convert_to_tensor(
        value=f_relative_tolerance, dtype=dtype, name='f_relative_tolerance')
    x_tolerance = tf.convert_to_tensor(
        value=x_tolerance, dtype=dtype, name='x_tolerance')
    max_iterations = tf.convert_to_tensor(
        value=max_iterations, name='max_iterations')

    # The `state` here is a `LBfgsOptimizerResults` tuple with values for the
    # current state of the algorithm computation.
    def _cond(state):
      """Continue if iterations remain and stopping condition is not met."""
      return ((state.num_iterations < max_iterations) &
              tf.logical_not(stopping_condition(state.converged, state.failed)))

    def _body(current_state):
      """Main optimization loop."""
      search_direction = _get_search_direction(current_state)

      # TODO(b/120134934): Check if the derivative at the start point is not
      # negative, if so then reset position/gradient deltas and recompute
      # search direction.

      next_state = bfgs_utils.line_search_step(
          current_state,
          value_and_gradients_function, search_direction,
          tolerance, f_relative_tolerance, x_tolerance, stopping_condition)

      # If not failed or converged, update the Hessian estimate.
      should_update = ~(next_state.converged | next_state.failed)
      state_after_inv_hessian_update = bfgs_utils.update_fields(
          next_state,
          position_deltas=_queue_push(
              current_state.position_deltas, should_update,
              next_state.position - current_state.position),
          gradient_deltas=_queue_push(
              current_state.gradient_deltas, should_update,
              next_state.objective_gradient - current_state.objective_gradient))
      return [state_after_inv_hessian_update]

    initial_state = _get_initial_state(value_and_gradients_function,
                                       initial_position,
                                       num_correction_pairs,
                                       tolerance)
    return tf.while_loop(
        cond=_cond,
        body=_body,
        loop_vars=[initial_state],
        parallel_iterations=parallel_iterations)[0]

Applies the L - BFGS algorithm to minimize a differentiable function .


def _get_initial_state(value_and_gradients_function,
                       initial_position,
                       num_correction_pairs,
                       tolerance):
  """Create LBfgsOptimizerResults with initial state of search procedure."""
  init_args = bfgs_utils.get_initial_state_args(
      value_and_gradients_function,
      initial_position,
      tolerance)
  empty_queue = _make_empty_queue_for(num_correction_pairs, initial_position)
  init_args.update(position_deltas=empty_queue, gradient_deltas=empty_queue)
  return LBfgsOptimizerResults(**init_args)

Create LBfgsOptimizerResults with initial state of search procedure .


def _get_search_direction(state):
  """Computes the search direction to follow at the current state.

  On the `k`-th iteration of the main L-BFGS algorithm, the state has collected
  the most recent `m` correction pairs in position_deltas and gradient_deltas,
  where `k = state.num_iterations` and `m = min(k, num_correction_pairs)`.

  Assuming these, the code below is an implementation of the L-BFGS two-loop
  recursion algorithm given by [Nocedal and Wright(2006)][1]:

  ```None
    q_direction = objective_gradient
    for i in reversed(range(m)):  # First loop.
      inv_rho[i] = gradient_deltas[i]^T * position_deltas[i]
      alpha[i] = position_deltas[i]^T * q_direction / inv_rho[i]
      q_direction = q_direction - alpha[i] * gradient_deltas[i]

    kth_inv_hessian_factor = (gradient_deltas[-1]^T * position_deltas[-1] /
                              gradient_deltas[-1]^T * gradient_deltas[-1])
    r_direction = kth_inv_hessian_factor * I * q_direction

    for i in range(m):  # Second loop.
      beta = gradient_deltas[i]^T * r_direction / inv_rho[i]
      r_direction = r_direction + position_deltas[i] * (alpha[i] - beta)

    return -r_direction  # Approximates - H_k * objective_gradient.
  ```

  Args:
    state: A `LBfgsOptimizerResults` tuple with the current state of the
      search procedure.

  Returns:
    A real `Tensor` of the same shape as the `state.position`. The direction
    along which to perform line search.
  """
  # The number of correction pairs that have been collected so far.
  num_elements = tf.minimum(
      state.num_iterations,
      distribution_util.prefer_static_shape(state.position_deltas)[0])

  def _two_loop_algorithm():
    """L-BFGS two-loop algorithm."""
    # Correction pairs are always appended to the end, so only the latest
    # `num_elements` vectors have valid position/gradient deltas.
    position_deltas = state.position_deltas[-num_elements:]
    gradient_deltas = state.gradient_deltas[-num_elements:]

    # Pre-compute all `inv_rho[i]`s.
    inv_rhos = tf.reduce_sum(
        input_tensor=gradient_deltas * position_deltas, axis=-1)

    def first_loop(acc, args):
      _, q_direction = acc
      position_delta, gradient_delta, inv_rho = args
      alpha = tf.reduce_sum(
          input_tensor=position_delta * q_direction, axis=-1) / inv_rho
      direction_delta = tf.expand_dims(alpha, axis=-1) * gradient_delta
      return (alpha, q_direction - direction_delta)

    # Run first loop body computing and collecting `alpha[i]`s, while also
    # computing the updated `q_direction` at each step.
    zero = tf.zeros_like(inv_rhos[0])
    alphas, q_directions = tf.scan(
        first_loop, [position_deltas, gradient_deltas, inv_rhos],
        initializer=(zero, state.objective_gradient), reverse=True)

    # We use `H^0_k = gamma_k * I` as an estimate for the initial inverse
    # hessian for the k-th iteration; then `r_direction = H^0_k * q_direction`.
    gamma_k = inv_rhos[-1] / tf.reduce_sum(
        input_tensor=gradient_deltas[-1] * gradient_deltas[-1], axis=-1)
    r_direction = tf.expand_dims(gamma_k, axis=-1) * q_directions[0]

    def second_loop(r_direction, args):
      alpha, position_delta, gradient_delta, inv_rho = args
      beta = tf.reduce_sum(
          input_tensor=gradient_delta * r_direction, axis=-1) / inv_rho
      direction_delta = tf.expand_dims(alpha - beta, axis=-1) * position_delta
      return r_direction + direction_delta

    # Finally, run second loop body computing the updated `r_direction` at each
    # step.
    r_directions = tf.scan(
        second_loop, [alphas, position_deltas, gradient_deltas, inv_rhos],
        initializer=r_direction)
    return -r_directions[-1]

  return prefer_static.cond(tf.equal(num_elements, 0),
                            (lambda: -state.objective_gradient),
                            _two_loop_algorithm)

Computes the search direction to follow at the current state .


def _make_empty_queue_for(k, element):
  """Creates a `tf.Tensor` suitable to hold `k` element-shaped tensors.

  For example:

  ```python
    element = tf.constant([[0., 1., 2., 3., 4.],
                           [5., 6., 7., 8., 9.]])

    # A queue capable of holding 3 elements.
    _make_empty_queue_for(3, element)
    # => [[[ 0.,  0.,  0.,  0.,  0.],
    #      [ 0.,  0.,  0.,  0.,  0.]],
    #
    #     [[ 0.,  0.,  0.,  0.,  0.],
    #      [ 0.,  0.,  0.,  0.,  0.]],
    #
    #     [[ 0.,  0.,  0.,  0.,  0.],
    #      [ 0.,  0.,  0.,  0.,  0.]]]
  ```

  Args:
    k: A positive scalar integer, number of elements that each queue will hold.
    element: A `tf.Tensor`, only its shape and dtype information are relevant.

  Returns:
    A zero-filed `tf.Tensor` of shape `(k,) + tf.shape(element)` and same dtype
    as `element`.
  """
  queue_shape = tf.concat(
      [[k], distribution_util.prefer_static_shape(element)], axis=0)
  return tf.zeros(queue_shape, dtype=element.dtype.base_dtype)

def _make_empty_queue_for ( k , element ) : queue_shape = tf . concat ( [ [ k ] , distribution_util . prefer_static_shape ( element ) ] , axis = 0 ) return tf . zeros ( queue_shape , dtype = element . dtype . base_dtype )

Creates a tf . Tensor suitable to hold k element - shaped tensors .


def _queue_push(queue, should_update, new_vecs):
  """Conditionally push new vectors into a batch of first-in-first-out queues.

  The `queue` of shape `[k, ..., n]` can be thought of as a batch of queues,
  each holding `k` n-D vectors; while `new_vecs` of shape `[..., n]` is a
  fresh new batch of n-D vectors. The `should_update` batch of Boolean scalars,
  i.e. shape `[...]`, indicates batch members whose corresponding n-D vector in
  `new_vecs` should be added at the back of its queue, pushing out the
  corresponding n-D vector from the front. Batch members in `new_vecs` for
  which `should_update` is False are ignored.

  Note: the choice of placing `k` at the dimension 0 of the queue is
  constrained by the L-BFGS two-loop algorithm above. The algorithm uses
  tf.scan to iterate over the `k` correction pairs simulatneously across all
  batches, and tf.scan itself can only iterate over dimension 0.

  For example:

  ```python
    k, b, n = (3, 2, 5)
    queue = tf.reshape(tf.range(30), (k, b, n))
    # => [[[ 0,  1,  2,  3,  4],
    #      [ 5,  6,  7,  8,  9]],
    #
    #     [[10, 11, 12, 13, 14],
    #      [15, 16, 17, 18, 19]],
    #
    #     [[20, 21, 22, 23, 24],
    #      [25, 26, 27, 28, 29]]]

    element = tf.reshape(tf.range(30, 40), (b, n))
    # => [[30, 31, 32, 33, 34],
          [35, 36, 37, 38, 39]]

    should_update = tf.constant([True, False])  # Shape: (b,)

    _queue_add(should_update, queue, element)
    # => [[[10, 11, 12, 13, 14],
    #      [ 5,  6,  7,  8,  9]],
    #
    #     [[20, 21, 22, 23, 24],
    #      [15, 16, 17, 18, 19]],
    #
    #     [[30, 31, 32, 33, 34],
    #      [25, 26, 27, 28, 29]]]
  ```

  Args:
    queue: A `tf.Tensor` of shape `[k, ..., n]`; a batch of queues each with
      `k` n-D vectors.
    should_update: A Boolean `tf.Tensor` of shape `[...]` indicating batch
      members where new vectors should be added to their queues.
    new_vecs: A `tf.Tensor` of shape `[..., n]`; a batch of n-D vectors to add
      at the end of their respective queues, pushing out the first element from
      each.

  Returns:
    A new `tf.Tensor` of shape `[k, ..., n]`.
  """
  new_queue = tf.concat([queue[1:], [new_vecs]], axis=0)
  update_pattern = tf.broadcast_to(
      should_update[tf.newaxis, ..., tf.newaxis],
      distribution_util.prefer_static_shape(queue))
  return tf.where(update_pattern, new_queue, queue)

def _queue_push ( queue , should_update , new_vecs ) : new_queue = tf . concat ( [ queue [ 1 : ] , [ new_vecs ] ] , axis = 0 ) update_pattern = tf . broadcast_to ( should_update [ tf . newaxis , ... , tf . newaxis ] , distribution_util . prefer_static_shape ( queue ) ) return tf . where ( update_pattern , new_queue , queue )

Conditionally push new vectors into a batch of first - in - first - out queues .


def _psd_mask(x):
  """Computes whether each square matrix in the input is positive semi-definite.

  Args:
    x: A floating-point `Tensor` of shape `[B1, ..., Bn, M, M]`.

  Returns:
    mask: A floating-point `Tensor` of shape `[B1, ... Bn]`.  Each
      scalar is 1 if the corresponding matrix was PSD, otherwise 0.
  """
  # Allegedly
  # https://scicomp.stackexchange.com/questions/12979/testing-if-a-matrix-is-positive-semi-definite
  # it is more efficient to test for positive semi-definiteness by
  # trying to compute the Cholesky decomposition -- the matrix is PSD
  # if you succeed and not PSD if you fail.  However, TensorFlow's
  # Cholesky raises an exception if _any_ of the input matrices are
  # not PSD, from which I don't know how to extract _which ones_, so I
  # proceed by explicitly computing all the eigenvalues and checking
  # whether they are all positive or not.
  #
  # Also, as was discussed in the answer, it is somewhat dangerous to
  # treat SPD-ness as binary in floating-point arithmetic. Cholesky
  # factorization can complete and 'look' like everything is fine
  # (e.g., O(1) entries and a diagonal of all ones) but the matrix can
  # have an exponential condition number.
  eigenvalues, _ = tf.linalg.eigh(x)
  return tf.cast(
      tf.reduce_min(input_tensor=eigenvalues, axis=-1) >= 0, dtype=x.dtype)

Computes whether each square matrix in the input is positive semi - definite .


def _det_large_enough_mask(x, det_bounds):
  """Returns whether the input matches the given determinant limit.

  Args:
    x: A floating-point `Tensor` of shape `[B1, ..., Bn, M, M]`.
    det_bounds: A floating-point `Tensor` that must broadcast to shape
      `[B1, ..., Bn]`, giving the desired lower bound on the
      determinants in `x`.

  Returns:
    mask: A floating-point `Tensor` of shape [B1, ..., Bn].  Each
      scalar is 1 if the corresponding matrix had determinant above
      the corresponding bound, otherwise 0.
  """
  # For the curious: I wonder whether it is possible and desirable to
  # use a Cholesky decomposition-based algorithm for this, since the
  # only matrices whose determinant this code cares about will be PSD.
  # Didn't figure out how to code that in TensorFlow.
  #
  # Expert opinion is that it would be about twice as fast since
  # Cholesky is roughly half the cost of Gaussian Elimination with
  # Partial Pivoting. But this is less of an impact than the switch in
  # _psd_mask.
  return tf.cast(tf.linalg.det(x) > det_bounds, dtype=x.dtype)

Returns whether the input matches the given determinant limit .


def _uniform_correlation_like_matrix(num_rows, batch_shape, dtype, seed):
  """Returns a uniformly random `Tensor` of "correlation-like" matrices.

  A "correlation-like" matrix is a symmetric square matrix with all entries
  between -1 and 1 (inclusive) and 1s on the main diagonal.  Of these,
  the ones that are positive semi-definite are exactly the correlation
  matrices.

  Args:
    num_rows: Python `int` dimension of the correlation-like matrices.
    batch_shape: `Tensor` or Python `tuple` of `int` shape of the
      batch to return.
    dtype: `dtype` of the `Tensor` to return.
    seed: Random seed.

  Returns:
    matrices: A `Tensor` of shape `batch_shape + [num_rows, num_rows]`
      and dtype `dtype`.  Each entry is in [-1, 1], and each matrix
      along the bottom two dimensions is symmetric and has 1s on the
      main diagonal.
  """
  num_entries = num_rows * (num_rows + 1) / 2
  ones = tf.ones(shape=[num_entries], dtype=dtype)
  # It seems wasteful to generate random values for the diagonal since
  # I am going to throw them away, but `fill_triangular` fills the
  # diagonal, so I probably need them.
  # It's not impossible that it would be more efficient to just fill
  # the whole matrix with random values instead of messing with
  # `fill_triangular`.  Then would need to filter almost half out with
  # `matrix_band_part`.
  unifs = uniform.Uniform(-ones, ones).sample(batch_shape, seed=seed)
  tril = util.fill_triangular(unifs)
  symmetric = tril + tf.linalg.matrix_transpose(tril)
  diagonal_ones = tf.ones(
      shape=util.pad(batch_shape, axis=0, back=True, value=num_rows),
      dtype=dtype)
  return tf.linalg.set_diag(symmetric, diagonal_ones)

Returns a uniformly random Tensor of correlation - like matrices .


def correlation_matrix_volume_rejection_samples(
    det_bounds, dim, sample_shape, dtype, seed):
  """Returns rejection samples from trying to get good correlation matrices.

  The proposal being rejected from is the uniform distribution on
  "correlation-like" matrices.  We say a matrix is "correlation-like"
  if it is a symmetric square matrix with all entries between -1 and 1
  (inclusive) and 1s on the main diagonal.  Of these, the ones that
  are positive semi-definite are exactly the correlation matrices.

  The rejection algorithm, then, is to sample a `Tensor` of
  `sample_shape` correlation-like matrices of dimensions `dim` by
  `dim`, and check each one for (i) being a correlation matrix (i.e.,
  PSD), and (ii) having determinant at least the corresponding entry
  of `det_bounds`.

  Args:
    det_bounds: A `Tensor` of lower bounds on the determinants of
      acceptable matrices.  The shape must broadcast with `sample_shape`.
    dim: A Python `int` dimension of correlation matrices to sample.
    sample_shape: Python `tuple` of `int` shape of the samples to
      compute, excluding the two matrix dimensions.
    dtype: The `dtype` in which to do the computation.
    seed: Random seed.

  Returns:
    weights: A `Tensor` of shape `sample_shape`.  Each entry is 0 if the
      corresponding matrix was not a correlation matrix, or had too
      small of a determinant.  Otherwise, the entry is the
      multiplicative inverse of the density of proposing that matrix
      uniformly, i.e., the volume of the set of `dim` by `dim`
      correlation-like matrices.
    volume: The volume of the set of `dim` by `dim` correlation-like
      matrices.
  """
  with tf.compat.v1.name_scope("rejection_sampler"):
    rej_proposals = _uniform_correlation_like_matrix(
        dim, sample_shape, dtype, seed=seed)
    rej_proposal_volume = 2. ** (dim * (dim - 1) / 2.)
    # The density of proposing any given point is 1 / rej_proposal_volume;
    # The weight of that point should be scaled by
    # 1 / density = rej_proposal_volume.
    rej_weights = rej_proposal_volume * _psd_mask(
        rej_proposals) * _det_large_enough_mask(rej_proposals, det_bounds)
    return rej_weights, rej_proposal_volume

Returns rejection samples from trying to get good correlation matrices .


def _clopper_pearson_confidence_interval(samples, error_rate):
  """Computes a confidence interval for the mean of the given 1-D distribution.

  Assumes (and checks) that the given distribution is Bernoulli, i.e.,
  takes only two values.  This licenses using the CDF of the binomial
  distribution for the confidence, which is tighter (for extreme
  probabilities) than the DKWM inequality.  The method is known as the
  [Clopper-Pearson method]
  (https://en.wikipedia.org/wiki/Binomial_proportion_confidence_interval).

  Assumes:

  - The given samples were drawn iid from the distribution of interest.

  - The given distribution is a Bernoulli, i.e., supported only on
    low and high.

  Guarantees:

  - The probability (over the randomness of drawing the given sample)
    that the true mean is outside the returned interval is no more
    than the given error_rate.

  Args:
    samples: `np.ndarray` of samples drawn iid from the distribution
      of interest.
    error_rate: Python `float` admissible rate of mistakes.

  Returns:
    low: Lower bound of confidence interval.
    high: Upper bound of confidence interval.

  Raises:
    ValueError: If `samples` has rank other than 1 (batch semantics
      are not implemented), or if `samples` contains values other than
      `low` or `high` (as that makes the distribution not Bernoulli).
  """
  # TODO(b/78025336) Migrate this confidence interval function
  # to statistical_testing.py.  In order to do that
  # - Get the binomial CDF from the Binomial distribution
  # - Implement scalar root finding in TF.  Batch bisection search
  #   shouldn't be too hard, and is definitely good enough for this
  #   problem.  Batching the Brent algorithm (from scipy) that is used
  #   here may be more involved, but may also not be necessary---it's
  #   only used here because scipy made it convenient.  In particular,
  #   robustness is more important than speed here, which may make
  #   bisection search actively better.
  # - The rest is just a matter of rewriting in the appropriate style.
  if optimize is None or stats is None:
    raise ValueError(
        "Scipy is required for computing Clopper-Pearson confidence intervals")
  if len(samples.shape) != 1:
    raise ValueError("Batch semantics not implemented")
  n = len(samples)
  low = np.amin(samples)
  high = np.amax(samples)
  successes = np.count_nonzero(samples - low)
  failures = np.count_nonzero(samples - high)
  if successes + failures != n:
    uniques = np.unique(samples)
    msg = ("Purportedly Bernoulli distribution had distinct samples"
           " {}, {}, and {}".format(uniques[0], uniques[1], uniques[2]))
    raise ValueError(msg)
  def p_small_enough(p):
    prob = stats.binom.logcdf(successes, n, p)
    return prob - np.log(error_rate / 2.)
  def p_big_enough(p):
    prob = stats.binom.logsf(successes, n, p)
    return prob - np.log(error_rate / 2.)
  high_p = optimize.brentq(
      p_small_enough, float(successes) / n, 1., rtol=1e-9)
  low_p = optimize.brentq(
      p_big_enough, 0., float(successes) / n, rtol=1e-9)
  low_interval = low + (high - low) * low_p
  high_interval = low + (high - low) * high_p
  return (low_interval, high_interval)

Computes a confidence interval for the mean of the given 1 - D distribution .


def compute_true_volumes(
    det_bounds, dim, num_samples, error_rate=1e-6, seed=42):
  """Returns confidence intervals for the desired correlation matrix volumes.

  The confidence intervals are computed by the [Clopper-Pearson method]
  (https://en.wikipedia.org/wiki/Binomial_proportion_confidence_interval).

  Args:
    det_bounds: A rank-1 numpy array of lower bounds on the
      determinants of acceptable matrices.  Entries must be unique.
    dim: A Python `int` dimension of correlation matrices to sample.
    num_samples: The number of samples to draw.
    error_rate: The statistical significance of the returned
      confidence intervals.  The significance is broadcast: Each
      returned interval separately may be incorrect with probability
      (under the sample of correlation-like matrices drawn internally)
      at most `error_rate`.
    seed: Random seed.

  Returns:
    bounds: A Python `dict` mapping each determinant bound to the low, high
      tuple giving the confidence interval.
  """
  bounds = {}
  with tf.compat.v1.Session() as sess:
    rej_weights, _ = correlation_matrix_volume_rejection_samples(
        det_bounds, dim, [num_samples, len(det_bounds)], np.float32, seed=seed)
    rej_weights = sess.run(rej_weights)
    for rw, det in zip(np.rollaxis(rej_weights, 1), det_bounds):
      template = ("Estimating volume of {}x{} correlation "
                  "matrices with determinant >= {}.")
      print(template.format(dim, dim, det))
      sys.stdout.flush()
      bounds[det] = _clopper_pearson_confidence_interval(
          rw, error_rate=error_rate)
    return bounds

Returns confidence intervals for the desired correlation matrix volumes .


def _von_mises_cdf_series(x, concentration, num_terms, dtype):
  """Computes the von Mises CDF and its derivative via series expansion."""
  # Keep the number of terms as a float. It should be a small integer, so
  # exactly representable as a float.
  num_terms = tf.cast(num_terms, dtype=dtype)

  def loop_body(n, rn, drn_dconcentration, vn, dvn_dconcentration):
    """One iteration of the series loop."""

    denominator = 2. * n / concentration + rn
    ddenominator_dk = -2. * n / concentration ** 2 + drn_dconcentration
    rn = 1. / denominator
    drn_dconcentration = -ddenominator_dk / denominator ** 2

    multiplier = tf.sin(n * x) / n + vn
    vn = rn * multiplier
    dvn_dconcentration = (drn_dconcentration * multiplier +
                          rn * dvn_dconcentration)
    n -= 1.

    return n, rn, drn_dconcentration, vn, dvn_dconcentration

  (_, _, _, vn, dvn_dconcentration) = tf.while_loop(
      cond=lambda n, *_: n > 0.,
      body=loop_body,
      loop_vars=(
          num_terms,  # n
          tf.zeros_like(x, name="rn"),
          tf.zeros_like(x, name="drn_dconcentration"),
          tf.zeros_like(x, name="vn"),
          tf.zeros_like(x, name="dvn_dconcentration"),
      ),
  )

  cdf = .5 + x / (2. * np.pi) + vn / np.pi
  dcdf_dconcentration = dvn_dconcentration / np.pi

  # Clip the result to [0, 1].
  cdf_clipped = tf.clip_by_value(cdf, 0., 1.)
  # The clipped values do not depend on concentration anymore, so set their
  # derivative to zero.
  dcdf_dconcentration *= tf.cast((cdf >= 0.) & (cdf <= 1.), dtype)

  return cdf_clipped, dcdf_dconcentration

Computes the von Mises CDF and its derivative via series expansion .


def _von_mises_cdf_normal(x, concentration, dtype):
  """Computes the von Mises CDF and its derivative via Normal approximation."""

  def cdf_func(concentration):
    """A helper function that is passed to value_and_gradient."""
    # z is an "almost Normally distributed" random variable.
    z = ((np.sqrt(2. / np.pi) / tf.math.bessel_i0e(concentration)) *
         tf.sin(.5 * x))

    # This is the correction described in [1] which reduces the error
    # of the Normal approximation.
    z2 = z ** 2
    z3 = z2 * z
    z4 = z2 ** 2
    c = 24. * concentration
    c1 = 56.

    xi = z - z3 / ((c - 2. * z2 - 16.) / 3. -
                   (z4 + (7. / 4.) * z2 + 167. / 2.) / (c - c1 - z2 + 3.)) ** 2

    distrib = normal.Normal(tf.cast(0., dtype), tf.cast(1., dtype))

    return distrib.cdf(xi)

  return value_and_gradient(cdf_func, concentration)

Computes the von Mises CDF and its derivative via Normal approximation .