def load_bernoulli_mnist_dataset(directory, split_name):
  """Returns Hugo Larochelle's binary static MNIST tf.data.Dataset."""
  amat_file = download(directory, FILE_TEMPLATE.format(split=split_name))
  dataset = tf.data.TextLineDataset(amat_file)
  str_to_arr = lambda string: np.array([c == b"1" for c in string.split()])

  def _parser(s):
    booltensor = tf.compat.v1.py_func(str_to_arr, [s], tf.bool)
    reshaped = tf.reshape(booltensor, [28, 28, 1])
    return tf.cast(reshaped, dtype=tf.float32), tf.constant(0, tf.int32)

  return dataset.map(_parser)

def load_bernoulli_mnist_dataset ( directory , split_name ) : amat_file = download ( directory , FILE_TEMPLATE . format ( split = split_name ) ) dataset = tf . data . TextLineDataset ( amat_file ) str_to_arr = lambda string : np . array ( [ c == b"1" for c in string . split ( ) ] ) def _parser ( s ) : booltensor = tf . compat . v1 . py_func ( str_to_arr , [ s ] , tf . bool ) reshaped = tf . reshape ( booltensor , [ 28 , 28 , 1 ] ) return tf . cast ( reshaped , dtype = tf . float32 ) , tf . constant ( 0 , tf . int32 ) return dataset . map ( _parser )

Returns Hugo Larochelle s binary static MNIST tf . data . Dataset .


def as_numpy_dtype(dtype):
  """Returns a `np.dtype` based on this `dtype`."""
  dtype = tf.as_dtype(dtype)
  if hasattr(dtype, 'as_numpy_dtype'):
    return dtype.as_numpy_dtype
  return dtype

def as_numpy_dtype ( dtype ) : dtype = tf . as_dtype ( dtype ) if hasattr ( dtype , 'as_numpy_dtype' ) : return dtype . as_numpy_dtype return dtype

Returns a np . dtype based on this dtype .


def base_dtype(dtype):
  """Returns a non-reference `dtype` based on this `dtype`."""
  dtype = tf.as_dtype(dtype)
  if hasattr(dtype, 'base_dtype'):
    return dtype.base_dtype
  return dtype

def base_dtype ( dtype ) : dtype = tf . as_dtype ( dtype ) if hasattr ( dtype , 'base_dtype' ) : return dtype . base_dtype return dtype

Returns a non - reference dtype based on this dtype .


def is_bool(dtype):
  """Returns whether this is a boolean data type."""
  dtype = tf.as_dtype(dtype)
  if hasattr(dtype, 'is_bool'):
    return dtype.is_bool
  # We use `kind` because:
  # np.issubdtype(np.uint8, np.bool) == True.
  return np.dtype(dtype).kind == 'b'

Returns whether this is a boolean data type .


def is_complex(dtype):
  """Returns whether this is a complex floating point type."""
  dtype = tf.as_dtype(dtype)
  if hasattr(dtype, 'is_complex'):
    return dtype.is_complex
  return np.issubdtype(np.dtype(dtype), np.complex)

def is_complex ( dtype ) : dtype = tf . as_dtype ( dtype ) if hasattr ( dtype , 'is_complex' ) : return dtype . is_complex return np . issubdtype ( np . dtype ( dtype ) , np . complex )

Returns whether this is a complex floating point type .


def max(dtype):  # pylint: disable=redefined-builtin
  """Returns the maximum representable value in this data type."""
  dtype = tf.as_dtype(dtype)
  if hasattr(dtype, 'max'):
    return dtype.max
  use_finfo = is_floating(dtype) or is_complex(dtype)
  return np.finfo(dtype).max if use_finfo else np.iinfo(dtype).max

def
max
(
dtype
)
:
# pylint: disable=redefined-builtin
dtype
=
tf
.
as_dtype
(
dtype
)
if
hasattr
(
dtype
,
'max'
)
:
return
dtype
.
max
use_finfo
=
is_floating
(
dtype
)
or
is_complex
(
dtype
)
return
np
.
finfo
(
dtype
)
.
max
if
use_finfo
else
np
.
iinfo
(
dtype
)
.
max

Returns the maximum representable value in this data type .


def name(dtype):
  """Returns the string name for this `dtype`."""
  dtype = tf.as_dtype(dtype)
  if hasattr(dtype, 'name'):
    return dtype.name
  if hasattr(dtype, '__name__'):
    return dtype.__name__
  return str(dtype)

def name ( dtype ) : dtype = tf . as_dtype ( dtype ) if hasattr ( dtype , 'name' ) : return dtype . name if hasattr ( dtype , '__name__' ) : return dtype . __name__ return str ( dtype )

Returns the string name for this dtype .


def size(dtype):
  """Returns the number of bytes to represent this `dtype`."""
  dtype = tf.as_dtype(dtype)
  if hasattr(dtype, 'size'):
    return dtype.size
  return np.dtype(dtype).itemsize

def size ( dtype ) : dtype = tf . as_dtype ( dtype ) if hasattr ( dtype , 'size' ) : return dtype . size return np . dtype ( dtype ) . itemsize

Returns the number of bytes to represent this dtype .


def _assert_same_base_type(items, expected_type=None):
  r"""Asserts all items are of the same base type.

  Args:
    items: List of graph items (e.g., `Variable`, `Tensor`, `SparseTensor`,
        `Operation`, or `IndexedSlices`). Can include `None` elements, which
        will be ignored.
    expected_type: Expected type. If not specified, assert all items are
        of the same base type.

  Returns:
    Validated type, or none if neither expected_type nor items provided.

  Raises:
    ValueError: If any types do not match.
  """
  original_expected_type = expected_type
  mismatch = False
  for item in items:
    if item is not None:
      item_type = base_dtype(item.dtype)
      if not expected_type:
        expected_type = item_type
      elif expected_type != item_type:
        mismatch = True
        break
  if mismatch:
    # Loop back through and build up an informative error message (this is very
    # slow, so we don't do it unless we found an error above).
    expected_type = original_expected_type
    original_item_str = None
    get_name = lambda x: x.name if hasattr(x, 'name') else str(x)
    for item in items:
      if item is not None:
        item_type = base_dtype(item.dtype)
        if not expected_type:
          expected_type = item_type
          original_item_str = get_name(item)
        elif expected_type != item_type:
          raise ValueError(
              '{}, type={}, must be of the same type ({}){}.'.format(
                  get_name(item),
                  item_type,
                  expected_type,
                  ((' as {}'.format(original_item_str))
                   if original_item_str else '')))
    return expected_type  # Should be unreachable
  else:
    return expected_type

r Asserts all items are of the same base type .


def assert_same_float_dtype(tensors=None, dtype=None):
  """Validate and return float type based on `tensors` and `dtype`.

  For ops such as matrix multiplication, inputs and weights must be of the
  same float type. This function validates that all `tensors` are the same type,
  validates that type is `dtype` (if supplied), and returns the type. Type must
  be a floating point type. If neither `tensors` nor `dtype` is supplied,
  the function will return `dtypes.float32`.

  Args:
    tensors: Tensors of input values. Can include `None` elements, which will
      be ignored.
    dtype: Expected type.

  Returns:
    Validated type.

  Raises:
    ValueError: if neither `tensors` nor `dtype` is supplied, or result is not
      float, or the common type of the inputs is not a floating point type.
  """
  if tensors:
    dtype = _assert_same_base_type(tensors, dtype)
  if not dtype:
    dtype = tf.float32
  elif not is_floating(dtype):
    raise ValueError('Expected floating point type, got {}.'.format(dtype))
  return dtype

def
assert_same_float_dtype
(
tensors
=
,
dtype
=
)
:
if
tensors
:
dtype
=
_assert_same_base_type
(
tensors
,
dtype
)
if
not
dtype
:
dtype
=
tf
.
float32
elif
not
is_floating
(
dtype
)
:
raise
ValueError
(
'Expected floating point type, got {}.'
.
format
(
dtype
)
)
return
dtype

Validate and return float type based on tensors and dtype .


def minimize(objective_function,
             initial_simplex=None,
             initial_vertex=None,
             step_sizes=None,
             objective_at_initial_simplex=None,
             objective_at_initial_vertex=None,
             batch_evaluate_objective=False,
             func_tolerance=1e-8,
             position_tolerance=1e-8,
             parallel_iterations=1,
             max_iterations=None,
             reflection=None,
             expansion=None,
             contraction=None,
             shrinkage=None,
             name=None):
  """Minimum of the objective function using the Nelder Mead simplex algorithm.

  Performs an unconstrained minimization of a (possibly non-smooth) function
  using the Nelder Mead simplex method. Nelder Mead method does not support
  univariate functions. Hence the dimensions of the domain must be 2 or greater.
  For details of the algorithm, see
  [Press, Teukolsky, Vetterling and Flannery(2007)][1].

  Points in the domain of the objective function may be represented as a
  `Tensor` of general shape but with rank at least 1. The algorithm proceeds
  by modifying a full rank simplex in the domain. The initial simplex may
  either be specified by the user or can be constructed using a single vertex
  supplied by the user. In the latter case, if `v0` is the supplied vertex,
  the simplex is the convex hull of the set:

  ```None
  S = {v0} + {v0 + step_i * e_i}
  ```

  Here `e_i` is a vector which is `1` along the `i`-th axis and zero elsewhere
  and `step_i` is a characteristic length scale along the `i`-th axis. If the
  step size is not supplied by the user, a unit step size is used in every axis.
  Alternately, a single step size may be specified which is used for every
  axis. The most flexible option is to supply a bespoke step size for every
  axis.

  ### Usage:

  The following example demonstrates the usage of the Nelder Mead minimzation
  on a two dimensional problem with the minimum located at a non-differentiable
  point.

  ```python
    # The objective function
    def sqrt_quadratic(x):
      return tf.sqrt(tf.reduce_sum(x ** 2, axis=-1))

    start = tf.constant([6.0, -21.0])  # Starting point for the search.
    optim_results = tfp.optimizer.nelder_mead_minimize(
        sqrt_quadratic, initial_vertex=start, func_tolerance=1e-8,
        batch_evaluate_objective=True)

    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, np.array([0.0, 0.0]),
                                 atol=1e-7)
      # Print out the total number of function evaluations it took.
      print ("Function evaluations: %d" % results.num_objective_evaluations)
  ```

  ### References:
  [1]: William Press, Saul Teukolsky, William Vetterling and Brian Flannery.
    Numerical Recipes in C++, third edition. pp. 502-507. (2007).
    http://numerical.recipes/cpppages/chap0sel.pdf

  [2]: Jeffrey Lagarias, James Reeds, Margaret Wright and Paul Wright.
    Convergence properties of the Nelder-Mead simplex method in low dimensions,
    Siam J. Optim., Vol 9, No. 1, pp. 112-147. (1998).
    http://www.math.kent.edu/~reichel/courses/Opt/reading.material.2/nelder.mead.pdf

  [3]: Fuchang Gao and Lixing Han. Implementing the Nelder-Mead simplex
    algorithm with adaptive parameters. Computational Optimization and
    Applications, Vol 51, Issue 1, pp 259-277. (2012).
    https://pdfs.semanticscholar.org/15b4/c4aa7437df4d032c6ee6ce98d6030dd627be.pdf

  Args:
    objective_function:  A Python callable that accepts a point as a
      real `Tensor` and returns a `Tensor` of real dtype containing
      the value of the function at that point. The function
      to be minimized. If `batch_evaluate_objective` is `True`, the callable
      may be evaluated on a `Tensor` of shape `[n+1] + s ` where `n` is
      the dimension of the problem and `s` is the shape of a single point
      in the domain (so `n` is the size of a `Tensor` representing a
      single point).
      In this case, the expected return value is a `Tensor` of shape `[n+1]`.
      Note that this method does not support univariate functions so the problem
      dimension `n` must be strictly greater than 1.
    initial_simplex: (Optional) `Tensor` of real dtype. The initial simplex to
      start the search. If supplied, should be a `Tensor` of shape `[n+1] + s`
      where `n` is the dimension of the problem and `s` is the shape of a
      single point in the domain. Each row (i.e. the `Tensor` with a given
      value of the first index) is interpreted as a vertex of a simplex and
      hence the rows must be affinely independent. If not supplied, an axes
      aligned simplex is constructed using the `initial_vertex` and
      `step_sizes`. Only one and at least one of `initial_simplex` and
      `initial_vertex` must be supplied.
    initial_vertex: (Optional) `Tensor` of real dtype and any shape that can
      be consumed by the `objective_function`. A single point in the domain that
      will be used to construct an axes aligned initial simplex.
    step_sizes: (Optional) `Tensor` of real dtype and shape broadcasting
      compatible with `initial_vertex`. Supplies the simplex scale along each
      axes. Only used if `initial_simplex` is not supplied. See description
      above for details on how step sizes and initial vertex are used to
      construct the initial simplex.
    objective_at_initial_simplex: (Optional) Rank `1` `Tensor` of real dtype
      of a rank `1` `Tensor`. The value of the objective function at the
      initial simplex. May be supplied only if `initial_simplex` is
      supplied. If not supplied, it will be computed.
    objective_at_initial_vertex: (Optional) Scalar `Tensor` of real dtype. The
      value of the objective function at the initial vertex. May be supplied
      only if the `initial_vertex` is also supplied.
    batch_evaluate_objective: (Optional) Python `bool`. If True, the objective
      function will be evaluated on all the vertices of the simplex packed
      into a single tensor. If False, the objective will be mapped across each
      vertex separately. Evaluating the objective function in a batch allows
      use of vectorization and should be preferred if the objective function
      allows it.
    func_tolerance: (Optional) Scalar `Tensor` of real dtype. The algorithm
      stops if the absolute difference between the largest and the smallest
      function value on the vertices of the simplex is below this number.
    position_tolerance: (Optional) Scalar `Tensor` of real dtype. The
      algorithm stops if the largest absolute difference between the
      coordinates of the vertices is below this threshold.
    parallel_iterations: (Optional) Positive integer. The number of iterations
      allowed to run in parallel.
    max_iterations: (Optional) Scalar positive `Tensor` of dtype `int32`.
      The maximum number of iterations allowed. If `None` then no limit is
      applied.
    reflection: (Optional) Positive Scalar `Tensor` of same dtype as
      `initial_vertex`. This parameter controls the scaling of the reflected
      vertex. See, [Press et al(2007)][1] for details. If not specified,
      uses the dimension dependent prescription of [Gao and Han(2012)][3].
    expansion: (Optional) Positive Scalar `Tensor` of same dtype as
      `initial_vertex`. Should be greater than `1` and `reflection`. This
      parameter controls the expanded scaling of a reflected vertex.
      See, [Press et al(2007)][1] for details. If not specified, uses the
      dimension dependent prescription of [Gao and Han(2012)][3].
    contraction: (Optional) Positive scalar `Tensor` of same dtype as
      `initial_vertex`. Must be between `0` and `1`. This parameter controls
      the contraction of the reflected vertex when the objective function at
      the reflected point fails to show sufficient decrease.
      See, [Press et al(2007)][1] for more details. If not specified, uses
      the dimension dependent prescription of [Gao and Han(2012][3].
    shrinkage: (Optional) Positive scalar `Tensor` of same dtype as
      `initial_vertex`. Must be between `0` and `1`. This parameter is the scale
      by which the simplex is shrunk around the best point when the other
      steps fail to produce improvements.
      See, [Press et al(2007)][1] for more details. If not specified, uses
      the dimension dependent prescription of [Gao and Han(2012][3].
    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.
      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.
      final_simplex: The last simplex constructed before stopping.
      final_objective_values: The objective function evaluated at the
        vertices of the final simplex.
      initial_simplex: The starting simplex.
      initial_objective_values: The objective function evaluated at the
        vertices of the initial simplex.
      num_iterations: The number of iterations of the main algorithm body.

  Raises:
    ValueError: If any of the following conditions hold
      1. If none or more than one of `initial_simplex` and `initial_vertex` are
        supplied.
      2. If `initial_simplex` and `step_sizes` are both specified.
  """
  with tf.compat.v1.name_scope(name, 'minimize', [
      initial_simplex, initial_vertex, step_sizes, objective_at_initial_simplex,
      objective_at_initial_vertex, func_tolerance, position_tolerance
  ]):
    (
        dim,
        _,
        simplex,
        objective_at_simplex,
        num_evaluations
    ) = _prepare_args(objective_function,
                      initial_simplex,
                      initial_vertex,
                      step_sizes,
                      objective_at_initial_simplex,
                      objective_at_initial_vertex,
                      batch_evaluate_objective)
    domain_dtype = simplex.dtype
    (
        reflection,
        expansion,
        contraction,
        shrinkage
    ) = _resolve_parameters(dim,
                            reflection,
                            expansion,
                            contraction,
                            shrinkage,
                            domain_dtype)

    closure_kwargs = dict(
        objective_function=objective_function,
        dim=dim,
        func_tolerance=func_tolerance,
        position_tolerance=position_tolerance,
        batch_evaluate_objective=batch_evaluate_objective,
        reflection=reflection,
        expansion=expansion,
        contraction=contraction,
        shrinkage=shrinkage)

    def _loop_body(_, iterations, simplex, objective_at_simplex,
                   num_evaluations):
      (
          converged,
          next_simplex,
          next_objective,
          evaluations
      ) = nelder_mead_one_step(simplex, objective_at_simplex, **closure_kwargs)

      return (converged, iterations + 1, next_simplex, next_objective,
              num_evaluations + evaluations)

    initial_args = (False, 0, simplex, objective_at_simplex,
                    num_evaluations)
    # Loop until either we have converged or if the max iterations are supplied
    # then until we have converged or exhausted the available iteration budget.
    def _is_converged(converged, num_iterations, *ignored_args):  # pylint:disable=unused-argument
      # It is important to ensure that not_converged is a tensor. If
      # converged is not a tensor but a Python bool, then the overloaded
      # op '~' acts as bitwise complement so ~True = -2 and ~False = -1.
      # In that case, the loop will never terminate.
      not_converged = tf.logical_not(converged)
      return (not_converged if max_iterations is None
              else (not_converged & (num_iterations < max_iterations)))

    (converged, num_iterations, final_simplex, final_objective_values,
     final_evaluations) = tf.while_loop(
         cond=_is_converged,
         body=_loop_body,
         loop_vars=initial_args,
         parallel_iterations=parallel_iterations)
    order = tf.argsort(
        final_objective_values, direction='ASCENDING', stable=True)
    best_index = order[0]
    # The explicit cast to Tensor below is done to avoid returning a mixture
    # of Python types and Tensors which cause problems with session.run.
    # In the eager mode, converged may remain a Python bool. Trying to evaluate
    # the whole tuple in one evaluate call will raise an exception because
    # of the presence of non-tensors. This is very annoying so we explicitly
    # cast those arguments to Tensors.
    return NelderMeadOptimizerResults(
        converged=tf.convert_to_tensor(value=converged),
        num_objective_evaluations=final_evaluations,
        position=final_simplex[best_index],
        objective_value=final_objective_values[best_index],
        final_simplex=final_simplex,
        final_objective_values=final_objective_values,
        num_iterations=tf.convert_to_tensor(value=num_iterations),
        initial_simplex=simplex,
        initial_objective_values=objective_at_simplex)

Minimum of the objective function using the Nelder Mead simplex algorithm .


def nelder_mead_one_step(current_simplex,
                         current_objective_values,
                         objective_function=None,
                         dim=None,
                         func_tolerance=None,
                         position_tolerance=None,
                         batch_evaluate_objective=False,
                         reflection=None,
                         expansion=None,
                         contraction=None,
                         shrinkage=None,
                         name=None):
  """A single iteration of the Nelder Mead algorithm."""
  with tf.compat.v1.name_scope(name, 'nelder_mead_one_step'):
    domain_dtype = current_simplex.dtype.base_dtype
    order = tf.argsort(
        current_objective_values, direction='ASCENDING', stable=True)
    (
        best_index,
        worst_index,
        second_worst_index
    ) = order[0], order[-1], order[-2]

    worst_vertex = current_simplex[worst_index]

    (
        best_objective_value,
        worst_objective_value,
        second_worst_objective_value
    ) = (
        current_objective_values[best_index],
        current_objective_values[worst_index],
        current_objective_values[second_worst_index]
    )

    # Compute the centroid of the face opposite the worst vertex.
    face_centroid = tf.reduce_sum(
        input_tensor=current_simplex, axis=0) - worst_vertex
    face_centroid /= tf.cast(dim, domain_dtype)

    # Reflect the worst vertex through the opposite face.
    reflected = face_centroid + reflection * (face_centroid - worst_vertex)
    objective_at_reflected = objective_function(reflected)

    num_evaluations = 1
    has_converged = _check_convergence(current_simplex,
                                       current_simplex[best_index],
                                       best_objective_value,
                                       worst_objective_value,
                                       func_tolerance,
                                       position_tolerance)
    def _converged_fn():
      return (True, current_simplex, current_objective_values, 0)
    case0 = has_converged, _converged_fn
    accept_reflected = (
        (objective_at_reflected < second_worst_objective_value) &
        (objective_at_reflected >= best_objective_value))
    accept_reflected_fn = _accept_reflected_fn(current_simplex,
                                               current_objective_values,
                                               worst_index,
                                               reflected,
                                               objective_at_reflected)
    case1 = accept_reflected, accept_reflected_fn
    do_expansion = objective_at_reflected < best_objective_value
    expansion_fn = _expansion_fn(objective_function,
                                 current_simplex,
                                 current_objective_values,
                                 worst_index,
                                 reflected,
                                 objective_at_reflected,
                                 face_centroid,
                                 expansion)
    case2 = do_expansion, expansion_fn
    do_outside_contraction = (
        (objective_at_reflected < worst_objective_value) &
        (objective_at_reflected >= second_worst_objective_value)
    )
    outside_contraction_fn = _outside_contraction_fn(
        objective_function,
        current_simplex,
        current_objective_values,
        face_centroid,
        best_index,
        worst_index,
        reflected,
        objective_at_reflected,
        contraction,
        shrinkage,
        batch_evaluate_objective)
    case3 = do_outside_contraction, outside_contraction_fn
    default_fn = _inside_contraction_fn(objective_function,
                                        current_simplex,
                                        current_objective_values,
                                        face_centroid,
                                        best_index,
                                        worst_index,
                                        worst_objective_value,
                                        contraction,
                                        shrinkage,
                                        batch_evaluate_objective)
    (
        converged,
        next_simplex,
        next_objective_at_simplex,
        case_evals) = prefer_static.case([case0, case1, case2, case3],
                                         default=default_fn, exclusive=False)
    next_simplex.set_shape(current_simplex.shape)
    next_objective_at_simplex.set_shape(current_objective_values.shape)
    return (
        converged,
        next_simplex,
        next_objective_at_simplex,
        num_evaluations + case_evals
    )

A single iteration of the Nelder Mead algorithm .


def _accept_reflected_fn(simplex,
                         objective_values,
                         worst_index,
                         reflected,
                         objective_at_reflected):
  """Creates the condition function pair for a reflection to be accepted."""
  def _replace_worst_with_reflected():
    next_simplex = _replace_at_index(simplex, worst_index, reflected)
    next_objective_values = _replace_at_index(objective_values, worst_index,
                                              objective_at_reflected)
    return False, next_simplex, next_objective_values, 0
  return _replace_worst_with_reflected

Creates the condition function pair for a reflection to be accepted .


def _expansion_fn(objective_function,
                  simplex,
                  objective_values,
                  worst_index,
                  reflected,
                  objective_at_reflected,
                  face_centroid,
                  expansion):
  """Creates the condition function pair for an expansion."""
  def _expand_and_maybe_replace():
    """Performs the expansion step."""
    expanded = face_centroid + expansion * (reflected - face_centroid)
    expanded_objective_value = objective_function(expanded)
    expanded_is_better = (expanded_objective_value <
                          objective_at_reflected)
    accept_expanded_fn = lambda: (expanded, expanded_objective_value)
    accept_reflected_fn = lambda: (reflected, objective_at_reflected)
    next_pt, next_objective_value = prefer_static.cond(
        expanded_is_better, accept_expanded_fn, accept_reflected_fn)
    next_simplex = _replace_at_index(simplex, worst_index, next_pt)
    next_objective_at_simplex = _replace_at_index(objective_values,
                                                  worst_index,
                                                  next_objective_value)
    return False, next_simplex, next_objective_at_simplex, 1
  return _expand_and_maybe_replace

def
_expansion_fn
(
objective_function
,
simplex
,
objective_values
,
worst_index
,
reflected
,
objective_at_reflected
,
face_centroid
,
expansion
)
:
def
_expand_and_maybe_replace
(
)
:
"""Performs the expansion step."""
expanded
=
face_centroid
+
expansion
*
(
reflected
face_centroid
)
expanded_objective_value
=
objective_function
(
expanded
)
expanded_is_better
=
(
expanded_objective_value
<
objective_at_reflected
)
accept_expanded_fn
=
lambda
:
(
expanded
,
expanded_objective_value
)
accept_reflected_fn
=
lambda
:
(
reflected
,
objective_at_reflected
)
next_pt
,
next_objective_value
=
prefer_static
.
cond
(
expanded_is_better
,
accept_expanded_fn
,
accept_reflected_fn
)
next_simplex
=
_replace_at_index
(
simplex
,
worst_index
,
next_pt
)
next_objective_at_simplex
=
_replace_at_index
(
objective_values
,
worst_index
,
next_objective_value
)
return
False
,
next_simplex
,
next_objective_at_simplex
,
1
return
_expand_and_maybe_replace

Creates the condition function pair for an expansion .


def _outside_contraction_fn(objective_function,
                            simplex,
                            objective_values,
                            face_centroid,
                            best_index,
                            worst_index,
                            reflected,
                            objective_at_reflected,
                            contraction,
                            shrinkage,
                            batch_evaluate_objective):
  """Creates the condition function pair for an outside contraction."""
  def _contraction():
    """Performs a contraction."""
    contracted = face_centroid + contraction * (reflected - face_centroid)
    objective_at_contracted = objective_function(contracted)
    is_contracted_acceptable = objective_at_contracted <= objective_at_reflected
    def _accept_contraction():
      next_simplex = _replace_at_index(simplex, worst_index, contracted)
      objective_at_next_simplex = _replace_at_index(
          objective_values,
          worst_index,
          objective_at_contracted)
      return (False,
              next_simplex,
              objective_at_next_simplex,
              1)

    def _reject_contraction():
      return _shrink_towards_best(objective_function,
                                  simplex,
                                  best_index,
                                  shrinkage,
                                  batch_evaluate_objective)

    return prefer_static.cond(is_contracted_acceptable,
                              _accept_contraction,
                              _reject_contraction)
  return _contraction

def
_outside_contraction_fn
(
objective_function
,
simplex
,
objective_values
,
face_centroid
,
best_index
,
worst_index
,
reflected
,
objective_at_reflected
,
contraction
,
shrinkage
,
batch_evaluate_objective
)
:
def
_contraction
(
)
:
"""Performs a contraction."""
contracted
=
face_centroid
+
contraction
*
(
reflected
face_centroid
)
objective_at_contracted
=
objective_function
(
contracted
)
is_contracted_acceptable
=
objective_at_contracted
<=
objective_at_reflected
def
_accept_contraction
(
)
:
next_simplex
=
_replace_at_index
(
simplex
,
worst_index
,
contracted
)
objective_at_next_simplex
=
_replace_at_index
(
objective_values
,
worst_index
,
objective_at_contracted
)
return
(
False
,
next_simplex
,
objective_at_next_simplex
,
1
)
def
_reject_contraction
(
)
:
return
_shrink_towards_best
(
objective_function
,
simplex
,
best_index
,
shrinkage
,
batch_evaluate_objective
)
return
prefer_static
.
cond
(
is_contracted_acceptable
,
_accept_contraction
,
_reject_contraction
)
return
_contraction

Creates the condition function pair for an outside contraction .


def _shrink_towards_best(objective_function,
                         simplex,
                         best_index,
                         shrinkage,
                         batch_evaluate_objective):
  """Shrinks the simplex around the best vertex."""

  # If the contraction step fails to improve the average objective enough,
  # the simplex is shrunk towards the best vertex.
  best_vertex = simplex[best_index]
  shrunk_simplex = best_vertex + shrinkage * (simplex - best_vertex)
  objective_at_shrunk_simplex, evals = _evaluate_objective_multiple(
      objective_function,
      shrunk_simplex,
      batch_evaluate_objective)
  return (False,
          shrunk_simplex,
          objective_at_shrunk_simplex,
          evals)

Shrinks the simplex around the best vertex .


def _replace_at_index(x, index, replacement):
  """Replaces an element at supplied index."""
  x_new = tf.concat([x[:index], tf.expand_dims(replacement, axis=0),
                     x[(index + 1):]], axis=0)
  return x_new

def _replace_at_index ( x , index , replacement ) : x_new = tf . concat ( [ x [ : index ] , tf . expand_dims ( replacement , axis = 0 ) , x [ ( index + 1 ) : ] ] , axis = 0 ) return x_new


def _check_convergence(simplex,
                       best_vertex,
                       best_objective,
                       worst_objective,
                       func_tolerance,
                       position_tolerance):
  """Returns True if the simplex has converged.

  If the simplex size is smaller than the `position_tolerance` or the variation
  of the function value over the vertices of the simplex is smaller than the
  `func_tolerance` return True else False.

  Args:
    simplex: `Tensor` of real dtype. The simplex to test for convergence. For
      more details, see the docstring for `initial_simplex` argument
      of `minimize`.
    best_vertex: `Tensor` of real dtype and rank one less than `simplex`. The
      vertex with the best (i.e. smallest) objective value.
    best_objective: Scalar `Tensor` of real dtype. The best (i.e. smallest)
      value of the objective function at a vertex.
    worst_objective: Scalar `Tensor` of same dtype as `best_objective`. The
      worst (i.e. largest) value of the objective function at a vertex.
    func_tolerance: Scalar positive `Tensor`. The tolerance for the variation
      of the objective function value over the simplex. If the variation over
      the simplex vertices is below this threshold, convergence is True.
    position_tolerance: Scalar positive `Tensor`. The algorithm stops if the
      lengths (under the supremum norm) of edges connecting to the best vertex
      are below this threshold.

  Returns:
    has_converged: A scalar boolean `Tensor` indicating whether the algorithm
      is deemed to have converged.
  """
  objective_convergence = tf.abs(worst_objective -
                                 best_objective) < func_tolerance
  simplex_degeneracy = tf.reduce_max(
      input_tensor=tf.abs(simplex - best_vertex)) < position_tolerance
  return objective_convergence | simplex_degeneracy

Returns True if the simplex has converged .


def _prepare_args(objective_function,
                  initial_simplex,
                  initial_vertex,
                  step_sizes,
                  objective_at_initial_simplex,
                  objective_at_initial_vertex,
                  batch_evaluate_objective):
  """Computes the initial simplex and the objective values at the simplex.

  Args:
    objective_function:  A Python callable that accepts a point as a
      real `Tensor` and returns a `Tensor` of real dtype containing
      the value of the function at that point. The function
      to be evaluated at the simplex. If `batch_evaluate_objective` is `True`,
      the callable may be evaluated on a `Tensor` of shape `[n+1] + s `
      where `n` is the dimension of the problem and `s` is the shape of a
      single point in the domain (so `n` is the size of a `Tensor`
      representing a single point).
      In this case, the expected return value is a `Tensor` of shape `[n+1]`.
    initial_simplex: None or `Tensor` of real dtype. The initial simplex to
      start the search. If supplied, should be a `Tensor` of shape `[n+1] + s`
      where `n` is the dimension of the problem and `s` is the shape of a
      single point in the domain. Each row (i.e. the `Tensor` with a given
      value of the first index) is interpreted as a vertex of a simplex and
      hence the rows must be affinely independent. If not supplied, an axes
      aligned simplex is constructed using the `initial_vertex` and
      `step_sizes`. Only one and at least one of `initial_simplex` and
      `initial_vertex` must be supplied.
    initial_vertex: None or `Tensor` of real dtype and any shape that can
      be consumed by the `objective_function`. A single point in the domain that
      will be used to construct an axes aligned initial simplex.
    step_sizes: None or `Tensor` of real dtype and shape broadcasting
      compatible with `initial_vertex`. Supplies the simplex scale along each
      axes. Only used if `initial_simplex` is not supplied. See the docstring
      of `minimize` for more details.
    objective_at_initial_simplex: None or rank `1` `Tensor` of real dtype.
      The value of the objective function at the initial simplex.
      May be supplied only if `initial_simplex` is
      supplied. If not supplied, it will be computed.
    objective_at_initial_vertex: None or scalar `Tensor` of real dtype. The
      value of the objective function at the initial vertex. May be supplied
      only if the `initial_vertex` is also supplied.
    batch_evaluate_objective: Python `bool`. If True, the objective function
      will be evaluated on all the vertices of the simplex packed into a
      single tensor. If False, the objective will be mapped across each
      vertex separately.

  Returns:
    prepared_args: A tuple containing the following elements:
      dimension: Scalar `Tensor` of `int32` dtype. The dimension of the problem
        as inferred from the supplied arguments.
      num_vertices: Scalar `Tensor` of `int32` dtype. The number of vertices
        in the simplex.
      simplex: A `Tensor` of same dtype as `initial_simplex`
        (or `initial_vertex`). The first component of the shape of the
        `Tensor` is `num_vertices` and each element represents a vertex of
        the simplex.
      objective_at_simplex: A `Tensor` of same dtype as the dtype of the
        return value of objective_function. The shape is a vector of size
        `num_vertices`. The objective function evaluated at the simplex.
      num_evaluations: An `int32` scalar `Tensor`. The number of points on
        which the objective function was evaluated.

  Raises:
    ValueError: If any of the following conditions hold
      1. If none or more than one of `initial_simplex` and `initial_vertex` are
        supplied.
      2. If `initial_simplex` and `step_sizes` are both specified.
  """
  if objective_at_initial_simplex is not None and initial_simplex is None:
    raise ValueError('`objective_at_initial_simplex` specified but the'
                     '`initial_simplex` was not.')

  if objective_at_initial_vertex is not None and initial_vertex is None:
    raise ValueError('`objective_at_initial_vertex` specified but the'
                     '`initial_vertex` was not.')

  # The full simplex was specified.
  if initial_simplex is not None:
    if initial_vertex is not None:
      raise ValueError('Both `initial_simplex` and `initial_vertex` specified.'
                       ' Only one of the two should be specified.')

    if step_sizes is not None:
      raise ValueError('`step_sizes` must not be specified when an'
                       ' `initial_simplex` has been specified.')
    return _prepare_args_with_initial_simplex(objective_function,
                                              initial_simplex,
                                              objective_at_initial_simplex,
                                              batch_evaluate_objective)

  if initial_vertex is None:
    raise ValueError('One of `initial_simplex` or `initial_vertex`'
                     ' must be supplied')

  if step_sizes is None:
    step_sizes = _default_step_sizes(initial_vertex)

  return _prepare_args_with_initial_vertex(objective_function,
                                           initial_vertex,
                                           step_sizes,
                                           objective_at_initial_vertex,
                                           batch_evaluate_objective)

Computes the initial simplex and the objective values at the simplex .


def _prepare_args_with_initial_simplex(objective_function,
                                       initial_simplex,
                                       objective_at_initial_simplex,
                                       batch_evaluate_objective):
  """Evaluates the objective function at the specified initial simplex."""
  initial_simplex = tf.convert_to_tensor(value=initial_simplex)

  # If d is the dimension of the problem, the number of vertices in the
  # simplex should be d+1. From this, we can infer the number of dimensions
  # as n - 1 where n is the number of vertices specified.
  num_vertices = tf.shape(input=initial_simplex)[0]
  dim = num_vertices - 1
  num_evaluations = 0

  if objective_at_initial_simplex is None:
    objective_at_initial_simplex, n_evals = _evaluate_objective_multiple(
        objective_function, initial_simplex, batch_evaluate_objective)
    num_evaluations += n_evals
  objective_at_initial_simplex = tf.convert_to_tensor(
      value=objective_at_initial_simplex)
  return (dim,
          num_vertices,
          initial_simplex,
          objective_at_initial_simplex,
          num_evaluations)

Evaluates the objective function at the specified initial simplex .


def _prepare_args_with_initial_vertex(objective_function,
                                      initial_vertex,
                                      step_sizes,
                                      objective_at_initial_vertex,
                                      batch_evaluate_objective):
  """Constructs a standard axes aligned simplex."""
  dim = tf.size(input=initial_vertex)
  num_vertices = dim + 1
  unit_vectors_along_axes = tf.reshape(
      tf.eye(dim, dim, dtype=initial_vertex.dtype.base_dtype),
      tf.concat([[dim], tf.shape(input=initial_vertex)], axis=0))

  # If step_sizes does not broadcast to initial_vertex, the multiplication
  # in the second term will fail.
  simplex_face = initial_vertex + step_sizes * unit_vectors_along_axes
  simplex = tf.concat([tf.expand_dims(initial_vertex, axis=0),
                       simplex_face], axis=0)
  num_evaluations = 0
  # Evaluate the objective function at the simplex vertices.
  if objective_at_initial_vertex is None:
    objective_at_initial_vertex = objective_function(initial_vertex)
    num_evaluations += 1

  objective_at_simplex_face, num_evals = _evaluate_objective_multiple(
      objective_function, simplex_face, batch_evaluate_objective)
  num_evaluations += num_evals

  objective_at_simplex = tf.concat(
      [
          tf.expand_dims(objective_at_initial_vertex, axis=0),
          objective_at_simplex_face
      ], axis=0)

  return (dim,
          num_vertices,
          simplex,
          objective_at_simplex,
          num_evaluations)

Constructs a standard axes aligned simplex .


def _evaluate_objective_multiple(objective_function, arg_batch,
                                 batch_evaluate_objective):
  """Evaluates the objective function on a batch of points.

  If `batch_evaluate_objective` is True, returns
  `objective function(arg_batch)` else it maps the `objective_function`
  across the `arg_batch`.

  Args:
    objective_function: A Python callable that accepts a single `Tensor` of
      rank 'R > 1' and any shape 's' and returns a scalar `Tensor` of real dtype
      containing the value of the function at that point. If
      `batch a `Tensor` of shape `[batch_size] + s ` where `batch_size` is the
      size of the batch of args. In this case, the expected return value is a
      `Tensor` of shape `[batch_size]`.
    arg_batch: A `Tensor` of real dtype. The batch of arguments at which to
      evaluate the `objective_function`. If `batch_evaluate_objective` is False,
      `arg_batch` will be unpacked along the zeroth axis and the
      `objective_function` will be applied to each element.
    batch_evaluate_objective: `bool`. Whether the `objective_function` can
      evaluate a batch of arguments at once.

  Returns:
    A tuple containing:
      objective_values: A `Tensor` of real dtype and shape `[batch_size]`.
        The value of the objective function evaluated at the supplied
        `arg_batch`.
      num_evaluations: An `int32` scalar `Tensor`containing the number of
        points on which the objective function was evaluated (i.e `batch_size`).
  """
  n_points = tf.shape(input=arg_batch)[0]
  if batch_evaluate_objective:
    return objective_function(arg_batch), n_points
  return tf.map_fn(objective_function, arg_batch), n_points

Evaluates the objective function on a batch of points .


def plot_weight_posteriors(names, qm_vals, qs_vals, fname):
  """Save a PNG plot with histograms of weight means and stddevs.

  Args:
    names: A Python `iterable` of `str` variable names.
    qm_vals: A Python `iterable`, the same length as `names`,
      whose elements are Numpy `array`s, of any shape, containing
      posterior means of weight varibles.
    qs_vals: A Python `iterable`, the same length as `names`,
      whose elements are Numpy `array`s, of any shape, containing
      posterior standard deviations of weight varibles.
    fname: Python `str` filename to save the plot to.
  """
  fig = figure.Figure(figsize=(6, 3))
  canvas = backend_agg.FigureCanvasAgg(fig)

  ax = fig.add_subplot(1, 2, 1)
  for n, qm in zip(names, qm_vals):
    sns.distplot(qm.flatten(), ax=ax, label=n)
  ax.set_title("weight means")
  ax.set_xlim([-1.5, 1.5])
  ax.legend()

  ax = fig.add_subplot(1, 2, 2)
  for n, qs in zip(names, qs_vals):
    sns.distplot(qs.flatten(), ax=ax)
  ax.set_title("weight stddevs")
  ax.set_xlim([0, 1.])

  fig.tight_layout()
  canvas.print_figure(fname, format="png")
  print("saved {}".format(fname))

def
plot_weight_posteriors
(
names
,
qm_vals
,
qs_vals
,
fname
)
:
fig
=
figure
.
Figure
(
figsize
=
(
6
,
3
)
)
canvas
=
backend_agg
.
FigureCanvasAgg
(
fig
)
ax
=
fig
.
add_subplot
(
1
,
2
,
1
)
for
n
,
qm
in
zip
(
names
,
qm_vals
)
:
sns
.
distplot
(
qm
.
flatten
(
)
,
ax
=
ax
,
label
=
n
)
ax
.
set_title
(
"weight means"
)
ax
.
set_xlim
(
[
1.5
,
1.5
]
)
ax
.
legend
(
)
ax
=
fig
.
add_subplot
(
1
,
2
,
2
)
for
n
,
qs
in
zip
(
names
,
qs_vals
)
:
sns
.
distplot
(
qs
.
flatten
(
)
,
ax
=
ax
)
ax
.
set_title
(
"weight stddevs"
)
ax
.
set_xlim
(
[
0
,
1.
]
)
fig
.
tight_layout
(
)
canvas
.
print_figure
(
fname
,
format
=
"png"
)
print
(
"saved {}"
.
format
(
fname
)
)

Save a PNG plot with histograms of weight means and stddevs .


def plot_heldout_prediction(input_vals, probs,
                            fname, n=10, title=""):
  """Save a PNG plot visualizing posterior uncertainty on heldout data.

  Args:
    input_vals: A `float`-like Numpy `array` of shape
      `[num_heldout] + IMAGE_SHAPE`, containing heldout input images.
    probs: A `float`-like Numpy array of shape `[num_monte_carlo,
      num_heldout, num_classes]` containing Monte Carlo samples of
      class probabilities for each heldout sample.
    fname: Python `str` filename to save the plot to.
    n: Python `int` number of datapoints to vizualize.
    title: Python `str` title for the plot.
  """
  fig = figure.Figure(figsize=(9, 3*n))
  canvas = backend_agg.FigureCanvasAgg(fig)
  for i in range(n):
    ax = fig.add_subplot(n, 3, 3*i + 1)
    ax.imshow(input_vals[i, :].reshape(IMAGE_SHAPE[:-1]), interpolation="None")

    ax = fig.add_subplot(n, 3, 3*i + 2)
    for prob_sample in probs:
      sns.barplot(np.arange(10), prob_sample[i, :], alpha=0.1, ax=ax)
      ax.set_ylim([0, 1])
    ax.set_title("posterior samples")

    ax = fig.add_subplot(n, 3, 3*i + 3)
    sns.barplot(np.arange(10), np.mean(probs[:, i, :], axis=0), ax=ax)
    ax.set_ylim([0, 1])
    ax.set_title("predictive probs")
  fig.suptitle(title)
  fig.tight_layout()

  canvas.print_figure(fname, format="png")
  print("saved {}".format(fname))

def
plot_heldout_prediction
(
input_vals
,
probs
,
fname
,
n
=
10
,
title
=
""
)
:
fig
=
figure
.
Figure
(
figsize
=
(
9
,
3
*
n
)
)
canvas
=
backend_agg
.
FigureCanvasAgg
(
fig
)
for
i
in
range
(
n
)
:
ax
=
fig
.
add_subplot
(
n
,
3
,
3
*
i
+
1
)
ax
.
imshow
(
input_vals
[
i
,
:
]
.
reshape
(
IMAGE_SHAPE
[
:
1
]
)
,
interpolation
=
"None"
)
ax
=
fig
.
add_subplot
(
n
,
3
,
3
*
i
+
2
)
for
prob_sample
in
probs
:
sns
.
barplot
(
np
.
arange
(
10
)
,
prob_sample
[
i
,
:
]
,
alpha
=
0.1
,
ax
=
ax
)
ax
.
set_ylim
(
[
0
,
1
]
)
ax
.
set_title
(
"posterior samples"
)
ax
=
fig
.
add_subplot
(
n
,
3
,
3
*
i
+
3
)
sns
.
barplot
(
np
.
arange
(
10
)
,
np
.
mean
(
probs
[
:
,
i
,
:
]
,
axis
=
0
)
,
ax
=
ax
)
ax
.
set_ylim
(
[
0
,
1
]
)
ax
.
set_title
(
"predictive probs"
)
fig
.
suptitle
(
title
)
fig
.
tight_layout
(
)
canvas
.
print_figure
(
fname
,
format
=
"png"
)
print
(
"saved {}"
.
format
(
fname
)
)

Save a PNG plot visualizing posterior uncertainty on heldout data .


def build_fake_data(num_examples=10):
  """Build fake MNIST-style data for unit testing."""

  class Dummy(object):
    pass

  num_examples = 10
  mnist_data = Dummy()
  mnist_data.train = Dummy()
  mnist_data.train.images = np.float32(np.random.randn(
      num_examples, *IMAGE_SHAPE))
  mnist_data.train.labels = np.int32(np.random.permutation(
      np.arange(num_examples)))
  mnist_data.train.num_examples = num_examples
  mnist_data.validation = Dummy()
  mnist_data.validation.images = np.float32(np.random.randn(
      num_examples, *IMAGE_SHAPE))
  mnist_data.validation.labels = np.int32(np.random.permutation(
      np.arange(num_examples)))
  mnist_data.validation.num_examples = num_examples
  return mnist_data

def build_fake_data ( num_examples = 10 ) : class Dummy ( object ) : pass num_examples = 10 mnist_data = Dummy ( ) mnist_data . train = Dummy ( ) mnist_data . train . images = np . float32 ( np . random . randn ( num_examples , * IMAGE_SHAPE ) ) mnist_data . train . labels = np . int32 ( np . random . permutation ( np . arange ( num_examples ) ) ) mnist_data . train . num_examples = num_examples mnist_data . validation = Dummy ( ) mnist_data . validation . images = np . float32 ( np . random . randn ( num_examples , * IMAGE_SHAPE ) ) mnist_data . validation . labels = np . int32 ( np . random . permutation ( np . arange ( num_examples ) ) ) mnist_data . validation . num_examples = num_examples return mnist_data

Build fake MNIST - style data for unit testing .


def get_config(self):
    """Returns initializer configuration as a JSON-serializable dict."""
    return {
        'initializers': [
            tf.compat.v2.initializers.serialize(
                tf.keras.initializers.get(init))
            for init in self.initializers
        ],
        'sizes': self.sizes,
        'validate_args': self.validate_args,
    }

def get_config ( self ) : return { 'initializers' : [ tf . compat . v2 . initializers . serialize ( tf . keras . initializers . get ( init ) ) for init in self . initializers ] , 'sizes' : self . sizes , 'validate_args' : self . validate_args , }

Returns initializer configuration as a JSON - serializable dict .


def from_config(cls, config):
    """Instantiates an initializer from a configuration dictionary."""
    return cls(**{
        'initializers': [tf.compat.v2.initializers.deserialize(init)
                         for init in config.get('initializers', [])],
        'sizes': config.get('sizes', []),
        'validate_args': config.get('validate_args', False),
    })

def from_config ( cls , config ) : return cls ( * * { 'initializers' : [ tf . compat . v2 . initializers . deserialize ( init ) for init in config . get ( 'initializers' , [ ] ) ] , 'sizes' : config . get ( 'sizes' , [ ] ) , 'validate_args' : config . get ( 'validate_args' , False ) , } )

Instantiates an initializer from a configuration dictionary .


def _matmul(a, b,
            transpose_a=False, transpose_b=False,
            adjoint_a=False, adjoint_b=False,
            a_is_sparse=False, b_is_sparse=False,
            name=None):  # pylint: disable=unused-argument
  """Numpy matmul wrapper."""
  if a_is_sparse or b_is_sparse:
    raise NotImplementedError('Numpy backend does not support sparse matmul.')
  if transpose_a or adjoint_a:
    a = _matrix_transpose(a, conjugate=adjoint_a)
  if transpose_b or adjoint_b:
    b = _matrix_transpose(b, conjugate=adjoint_b)
  return np.matmul(a, b)


def _std_var_helper(self, statistic, statistic_name, statistic_ndims,
                      df_factor_fn):
    """Helper to compute stddev, covariance and variance."""
    df = tf.reshape(
        self.df,
        tf.concat([
            tf.shape(input=self.df),
            tf.ones([statistic_ndims], dtype=tf.int32)
        ], -1))
    df = _broadcast_to_shape(df, tf.shape(input=statistic))
    # We need to put the tf.where inside the outer tf.where to ensure we never
    # hit a NaN in the gradient.
    denom = tf.where(df > 2., df - 2., tf.ones_like(df))
    statistic = statistic * df_factor_fn(df / denom)
    # When 1 < df <= 2, stddev/variance are infinite.
    inf = dtype_util.as_numpy_dtype(self.dtype)(np.inf)
    result_where_defined = tf.where(
        df > 2., statistic, tf.fill(tf.shape(input=statistic), inf, name="inf"))

    if self.allow_nan_stats:
      nan = dtype_util.as_numpy_dtype(self.dtype)(np.nan)
      return tf.where(df > 1., result_where_defined,
                      tf.fill(tf.shape(input=statistic), nan, name="nan"))
    else:
      with tf.control_dependencies([
          assert_util.assert_less(
              tf.cast(1., self.dtype),
              df,
              message=statistic_name +
              " not defined for components of df <= 1"),
      ]):
        return tf.identity(result_where_defined)

Helper to compute stddev covariance and variance .


def assign_log_moving_mean_exp(
    log_mean_exp_var, log_value, decay, name=None):
  """Compute the log of the exponentially weighted moving mean of the exp.

  If `log_value` is a draw from a stationary random variable, this function
  approximates `log(E[exp(log_value)])`, i.e., a weighted log-sum-exp. More
  precisely, a `tf.Variable`, `log_mean_exp_var`, is updated by `log_value`
  using the following identity:

  ```none
  log_mean_exp_var =
  = log(decay exp(log_mean_exp_var) + (1 - decay) exp(log_value))
  = log(exp(log_mean_exp_var + log(decay)) + exp(log_value + log1p(-decay)))
  = log_mean_exp_var
    + log(  exp(log_mean_exp_var   - log_mean_exp_var + log(decay))
          + exp(log_value - log_mean_exp_var + log1p(-decay)))
  = log_mean_exp_var
    + log_sum_exp([log(decay), log_value - log_mean_exp_var + log1p(-decay)]).
  ```

  In addition to numerical stability, this formulation is advantageous because
  `log_mean_exp_var` can be updated in a lock-free manner, i.e., using
  `assign_add`. (Note: the updates are not thread-safe; it's just that the
  update to the tf.Variable is presumed efficient due to being lock-free.)

  Args:
    log_mean_exp_var: `float`-like `Variable` representing the log of the
      exponentially weighted moving mean of the exp. Same shape as `log_value`.
    log_value: `float`-like `Tensor` representing a new (streaming) observation.
      Same shape as `log_mean_exp_var`.
    decay: A `float`-like `Tensor`. The moving mean decay. Typically close to
      `1.`, e.g., `0.999`.
    name: Optional name of the returned operation.

  Returns:
    log_mean_exp_var: A reference to the input 'Variable' tensor with the
      `log_value`-updated log of the exponentially weighted moving mean of exp.

  Raises:
    TypeError: if `log_mean_exp_var` does not have float type `dtype`.
    TypeError: if `log_mean_exp_var`, `log_value`, `decay` have different
      `base_dtype`.
  """
  with tf.compat.v1.name_scope(name, "assign_log_moving_mean_exp",
                               [log_mean_exp_var, log_value, decay]):
    # We want to update the variable in a numerically stable and lock-free way.
    # To do this, observe that variable `x` updated by `v` is:
    # x = log(w exp(x) + (1-w) exp(v))
    #   = log(exp(x + log(w)) + exp(v + log1p(-w)))
    #   = x + log(exp(x - x + log(w)) + exp(v - x + log1p(-w)))
    #   = x + lse([log(w), v - x + log1p(-w)])
    with tf.compat.v1.colocate_with(log_mean_exp_var):
      base_dtype = log_mean_exp_var.dtype.base_dtype
      if not base_dtype.is_floating:
        raise TypeError(
            "log_mean_exp_var.base_dtype({}) does not have float type "
            "`dtype`.".format(base_dtype.name))
      log_value = tf.convert_to_tensor(
          value=log_value, dtype=base_dtype, name="log_value")
      decay = tf.convert_to_tensor(value=decay, dtype=base_dtype, name="decay")
      delta = (log_value - log_mean_exp_var)[tf.newaxis, ...]
      x = tf.concat([
          tf.math.log(decay) * tf.ones_like(delta),
          delta + tf.math.log1p(-decay)
      ],
                    axis=0)
      x = tf.reduce_logsumexp(input_tensor=x, axis=0)
      return log_mean_exp_var.assign_add(x)

Compute the log of the exponentially weighted moving mean of the exp .


def _make_columnar(self, x):
    """Ensures non-scalar input has at least one column.

    Example:
      If `x = [1, 2, 3]` then the output is `[[1], [2], [3]]`.

      If `x = [[1, 2, 3], [4, 5, 6]]` then the output is unchanged.

      If `x = 1` then the output is unchanged.

    Args:
      x: `Tensor`.

    Returns:
      columnar_x: `Tensor` with at least two dimensions.
    """
    if tensorshape_util.rank(x.shape) is not None:
      if tensorshape_util.rank(x.shape) == 1:
        x = x[tf.newaxis, :]
      return x
    shape = tf.shape(input=x)
    maybe_expanded_shape = tf.concat([
        shape[:-1],
        distribution_util.pick_vector(
            tf.equal(tf.rank(x), 1), [1], np.array([], dtype=np.int32)),
        shape[-1:],
    ], 0)
    return tf.reshape(x, maybe_expanded_shape)

def _make_columnar ( self , x ) : if tensorshape_util . rank ( x . shape ) is not None : if tensorshape_util . rank ( x . shape ) == 1 : x = x [ tf . newaxis , : ] return x shape = tf . shape ( input = x ) maybe_expanded_shape = tf . concat ( [ shape [ : - 1 ] , distribution_util . pick_vector ( tf . equal ( tf . rank ( x ) , 1 ) , [ 1 ] , np . array ( [ ] , dtype = np . int32 ) ) , shape [ - 1 : ] , ] , 0 ) return tf . reshape ( x , maybe_expanded_shape )

Ensures non - scalar input has at least one column .


def random_rademacher(shape, dtype=tf.float32, seed=None, name=None):
  """Generates `Tensor` consisting of `-1` or `+1`, chosen uniformly at random.

  For more details, see [Rademacher distribution](
  https://en.wikipedia.org/wiki/Rademacher_distribution).

  Args:
    shape: Vector-shaped, `int` `Tensor` representing shape of output.
    dtype: (Optional) TF `dtype` representing `dtype` of output.
    seed: (Optional) Python integer to seed the random number generator.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., 'random_rademacher').

  Returns:
    rademacher: `Tensor` with specified `shape` and `dtype` consisting of `-1`
      or `+1` chosen uniformly-at-random.
  """
  with tf.compat.v1.name_scope(name, 'random_rademacher', [shape, seed]):
    # Choose the dtype to cause `2 * random_bernoulli - 1` to run in the same
    # memory (host or device) as the downstream cast will want to put it.  The
    # convention on GPU is that int32 are in host memory and int64 are in device
    # memory.
    generation_dtype = tf.int64 if tf.as_dtype(dtype) != tf.int32 else tf.int32
    random_bernoulli = tf.random.uniform(
        shape, minval=0, maxval=2, dtype=generation_dtype, seed=seed)
    return tf.cast(2 * random_bernoulli - 1, dtype)

Generates Tensor consisting of - 1 or + 1 chosen uniformly at random .


def random_rayleigh(shape, scale=None, dtype=tf.float32, seed=None, name=None):
  """Generates `Tensor` of positive reals drawn from a Rayleigh distributions.

  The probability density function of a Rayleigh distribution with `scale`
  parameter is given by:

  ```none
  f(x) = x scale**-2 exp(-x**2 0.5 scale**-2)
  ```

  For more details, see [Rayleigh distribution](
  https://en.wikipedia.org/wiki/Rayleigh_distribution)

  Args:
    shape: Vector-shaped, `int` `Tensor` representing shape of output.
    scale: (Optional) Positive `float` `Tensor` representing `Rayleigh` scale.
      Default value: `None` (i.e., `scale = 1.`).
    dtype: (Optional) TF `dtype` representing `dtype` of output.
      Default value: `tf.float32`.
    seed: (Optional) Python integer to seed the random number generator.
      Default value: `None` (i.e., no seed).
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., 'random_rayleigh').

  Returns:
    rayleigh: `Tensor` with specified `shape` and `dtype` consisting of positive
      real values drawn from a Rayleigh distribution with specified `scale`.
  """
  with tf.compat.v1.name_scope(name, 'random_rayleigh', [shape, scale, seed]):
    if scale is not None:
      # Its important to expand the shape to match scale's, otherwise we won't
      # have independent draws.
      scale = tf.convert_to_tensor(value=scale, dtype=dtype, name='scale')
      shape = tf.broadcast_dynamic_shape(shape, tf.shape(input=scale))
    x = tf.sqrt(-2. * tf.math.log(
        tf.random.uniform(shape, minval=0, maxval=1, dtype=dtype, seed=seed)))
    if scale is None:
      return x
    return x * scale

Generates Tensor of positive reals drawn from a Rayleigh distributions .


def _pick_scalar_condition(pred, cond_true, cond_false):
  """Convenience function which chooses the condition based on the predicate."""
  # Note: This function is only valid if all of pred, cond_true, and cond_false
  # are scalars. This means its semantics are arguably more like tf.cond than
  # tf.where even though we use tf.where to implement it.
  pred_ = tf.get_static_value(tf.convert_to_tensor(value=pred))
  if pred_ is None:
    return tf.where(pred, cond_true, cond_false)
  return cond_true if pred_ else cond_false

Convenience function which chooses the condition based on the predicate .


def _finish_log_prob_for_one_fiber(self, y, x, ildj, event_ndims,
                                     **distribution_kwargs):
    """Finish computation of log_prob on one element of the inverse image."""
    x = self._maybe_rotate_dims(x, rotate_right=True)
    log_prob = self.distribution.log_prob(x, **distribution_kwargs)
    if self._is_maybe_event_override:
      log_prob = tf.reduce_sum(
          input_tensor=log_prob, axis=self._reduce_event_indices)
    log_prob += tf.cast(ildj, log_prob.dtype)
    if self._is_maybe_event_override and isinstance(event_ndims, int):
      tensorshape_util.set_shape(
          log_prob,
          tf.broadcast_static_shape(
              tensorshape_util.with_rank_at_least(y.shape, 1)[:-event_ndims],
              self.batch_shape))
    return log_prob

def _finish_log_prob_for_one_fiber ( self , y , x , ildj , event_ndims , * * distribution_kwargs ) : x = self . _maybe_rotate_dims ( x , rotate_right = True ) log_prob = self . distribution . log_prob ( x , * * distribution_kwargs ) if self . _is_maybe_event_override : log_prob = tf . reduce_sum ( input_tensor = log_prob , axis = self . _reduce_event_indices ) log_prob += tf . cast ( ildj , log_prob . dtype ) if self . _is_maybe_event_override and isinstance ( event_ndims , int ) : tensorshape_util . set_shape ( log_prob , tf . broadcast_static_shape ( tensorshape_util . with_rank_at_least ( y . shape , 1 ) [ : - event_ndims ] , self . batch_shape ) ) return log_prob

Finish computation of log_prob on one element of the inverse image .


def _finish_prob_for_one_fiber(self, y, x, ildj, event_ndims,
                                 **distribution_kwargs):
    """Finish computation of prob on one element of the inverse image."""
    x = self._maybe_rotate_dims(x, rotate_right=True)
    prob = self.distribution.prob(x, **distribution_kwargs)
    if self._is_maybe_event_override:
      prob = tf.reduce_prod(input_tensor=prob, axis=self._reduce_event_indices)
    prob *= tf.exp(tf.cast(ildj, prob.dtype))
    if self._is_maybe_event_override and isinstance(event_ndims, int):
      tensorshape_util.set_shape(
          prob,
          tf.broadcast_static_shape(
              tensorshape_util.with_rank_at_least(y.shape, 1)[:-event_ndims],
              self.batch_shape))
    return prob

def _finish_prob_for_one_fiber ( self , y , x , ildj , event_ndims , * * distribution_kwargs ) : x = self . _maybe_rotate_dims ( x , rotate_right = True ) prob = self . distribution . prob ( x , * * distribution_kwargs ) if self . _is_maybe_event_override : prob = tf . reduce_prod ( input_tensor = prob , axis = self . _reduce_event_indices ) prob *= tf . exp ( tf . cast ( ildj , prob . dtype ) ) if self . _is_maybe_event_override and isinstance ( event_ndims , int ) : tensorshape_util . set_shape ( prob , tf . broadcast_static_shape ( tensorshape_util . with_rank_at_least ( y . shape , 1 ) [ : - event_ndims ] , self . batch_shape ) ) return prob

Finish computation of prob on one element of the inverse image .


def _maybe_rotate_dims(self, x, rotate_right=False):
    """Helper which rolls left event_dims left or right event_dims right."""
    needs_rotation_const = tf.get_static_value(self._needs_rotation)
    if needs_rotation_const is not None and not needs_rotation_const:
      return x
    ndims = prefer_static.rank(x)
    n = (ndims - self._rotate_ndims) if rotate_right else self._rotate_ndims
    perm = prefer_static.concat([
        prefer_static.range(n, ndims), prefer_static.range(0, n)], axis=0)
    return tf.transpose(a=x, perm=perm)

def _maybe_rotate_dims ( self , x , rotate_right = False ) : needs_rotation_const = tf . get_static_value ( self . _needs_rotation ) if needs_rotation_const is not None and not needs_rotation_const : return x ndims = prefer_static . rank ( x ) n = ( ndims - self . _rotate_ndims ) if rotate_right else self . _rotate_ndims perm = prefer_static . concat ( [ prefer_static . range ( n , ndims ) , prefer_static . range ( 0 , n ) ] , axis = 0 ) return tf . transpose ( a = x , perm = perm )

Helper which rolls left event_dims left or right event_dims right .


def _undo_batch_normalization(x,
                              mean,
                              variance,
                              offset,
                              scale,
                              variance_epsilon,
                              name=None):
  r"""Inverse of tf.nn.batch_normalization.

  Args:
    x: Input `Tensor` of arbitrary dimensionality.
    mean: A mean `Tensor`.
    variance: A variance `Tensor`.
    offset: An offset `Tensor`, often denoted `beta` in equations, or
      None. If present, will be added to the normalized tensor.
    scale: A scale `Tensor`, often denoted `gamma` in equations, or
      `None`. If present, the scale is applied to the normalized tensor.
    variance_epsilon: A small `float` added to the minibatch `variance` to
      prevent dividing by zero.
    name: A name for this operation (optional).

  Returns:
    batch_unnormalized: The de-normalized, de-scaled, de-offset `Tensor`.
  """
  with tf.compat.v2.name_scope(name or "undo_batchnorm"):
    # inv = tf.rsqrt(variance + variance_epsilon)
    # if scale is not None:
    #   inv *= scale
    # return x * inv + (
    #     offset - mean * inv if offset is not None else -mean * inv)
    rescale = tf.sqrt(variance + variance_epsilon)
    if scale is not None:
      rescale /= scale
    batch_unnormalized = x * rescale + (
        mean - offset * rescale if offset is not None else mean)
    return batch_unnormalized

r Inverse of tf . nn . batch_normalization .


def _validate_bn_layer(self, layer):
    """Check for valid BatchNormalization layer.

    Args:
      layer: Instance of `tf.layers.BatchNormalization`.
    Raises:
      ValueError: If batchnorm_layer argument is not an instance of
      `tf.layers.BatchNormalization`, or if `batchnorm_layer.renorm=True` or
      if `batchnorm_layer.virtual_batch_size` is specified.
    """
    if (not isinstance(layer, tf.keras.layers.BatchNormalization) and
        not isinstance(layer, tf.compat.v1.layers.BatchNormalization)):
      raise ValueError(
          "batchnorm_layer must be an instance of BatchNormalization layer.")
    if layer.renorm:
      raise ValueError("BatchNorm Bijector does not support renormalization.")
    if layer.virtual_batch_size:
      raise ValueError(
          "BatchNorm Bijector does not support virtual batch sizes.")

Check for valid BatchNormalization layer .


def _slice_single_param(param, param_event_ndims, slices, dist_batch_shape):
  """Slices a single parameter of a distribution.

  Args:
    param: A `Tensor`, the original parameter to slice.
    param_event_ndims: `int` event parameterization rank for this parameter.
    slices: A `tuple` of normalized slices.
    dist_batch_shape: The distribution's batch shape `Tensor`.

  Returns:
    new_param: A `Tensor`, batch-sliced according to slices.
  """
  # Extend param shape with ones on the left to match dist_batch_shape.
  param_shape = tf.shape(input=param)
  insert_ones = tf.ones(
      [tf.size(input=dist_batch_shape) + param_event_ndims - tf.rank(param)],
      dtype=param_shape.dtype)
  new_param_shape = tf.concat([insert_ones, param_shape], axis=0)
  full_batch_param = tf.reshape(param, new_param_shape)
  param_slices = []
  # We separately track the batch axis from the parameter axis because we want
  # them to align for positive indexing, and be offset by param_event_ndims for
  # negative indexing.
  param_dim_idx = 0
  batch_dim_idx = 0
  for slc in slices:
    if slc is tf.newaxis:
      param_slices.append(slc)
      continue
    if slc is Ellipsis:
      if batch_dim_idx < 0:
        raise ValueError('Found multiple `...` in slices {}'.format(slices))
      param_slices.append(slc)
      # Switch over to negative indexing for the broadcast check.
      num_remaining_non_newaxis_slices = sum(
          [s is not tf.newaxis for s in slices[slices.index(Ellipsis) + 1:]])
      batch_dim_idx = -num_remaining_non_newaxis_slices
      param_dim_idx = batch_dim_idx - param_event_ndims
      continue
    # Find the batch dimension sizes for both parameter and distribution.
    param_dim_size = new_param_shape[param_dim_idx]
    batch_dim_size = dist_batch_shape[batch_dim_idx]
    is_broadcast = batch_dim_size > param_dim_size
    # Slices are denoted by start:stop:step.
    if isinstance(slc, slice):
      start, stop, step = slc.start, slc.stop, slc.step
      if start is not None:
        start = tf.where(is_broadcast, 0, start)
      if stop is not None:
        stop = tf.where(is_broadcast, 1, stop)
      if step is not None:
        step = tf.where(is_broadcast, 1, step)
      param_slices.append(slice(start, stop, step))
    else:  # int, or int Tensor, e.g. d[d.batch_shape_tensor()[0] // 2]
      param_slices.append(tf.where(is_broadcast, 0, slc))
    param_dim_idx += 1
    batch_dim_idx += 1
  param_slices.extend([ALL_SLICE] * param_event_ndims)
  return full_batch_param.__getitem__(param_slices)

Slices a single parameter of a distribution .


def _slice_params_to_dict(dist, params_event_ndims, slices):
  """Computes the override dictionary of sliced parameters.

  Args:
    dist: The tfd.Distribution being batch-sliced.
    params_event_ndims: Per-event parameter ranks, a `str->int` `dict`.
    slices: Slices as received by __getitem__.

  Returns:
    overrides: `str->Tensor` `dict` of batch-sliced parameter overrides.
  """
  override_dict = {}
  for param_name, param_event_ndims in six.iteritems(params_event_ndims):
    # Verify that either None or a legit value is in the parameters dict.
    if param_name not in dist.parameters:
      raise ValueError('Distribution {} is missing advertised '
                       'parameter {}'.format(dist, param_name))
    param = dist.parameters[param_name]
    if param is None:
      # some distributions have multiple possible parameterizations; this
      # param was not provided
      continue
    dtype = None
    if hasattr(dist, param_name):
      attr = getattr(dist, param_name)
      dtype = getattr(attr, 'dtype', None)
    if dtype is None:
      dtype = dist.dtype
      warnings.warn('Unable to find property getter for parameter Tensor {} '
                    'on {}, falling back to Distribution.dtype {}'.format(
                        param_name, dist, dtype))
    param = tf.convert_to_tensor(value=param, dtype=dtype)
    override_dict[param_name] = _slice_single_param(param, param_event_ndims,
                                                    slices,
                                                    dist.batch_shape_tensor())
  return override_dict

Computes the override dictionary of sliced parameters .


def _apply_single_step(dist, params_event_ndims, slices, params_overrides):
  """Applies a single slicing step to `dist`, returning a new instance."""
  if len(slices) == 1 and slices[0] == Ellipsis:
    # The path used by Distribution.copy: batch_slice(...args..., Ellipsis)
    override_dict = {}
  else:
    override_dict = _slice_params_to_dict(dist, params_event_ndims, slices)
  override_dict.update(params_overrides)
  parameters = dict(dist.parameters, **override_dict)
  new_dist = type(dist)(**parameters)
  return new_dist

Applies a single slicing step to dist returning a new instance .


def _apply_slice_sequence(dist, params_event_ndims, slice_overrides_seq):
  """Applies a sequence of slice or copy-with-overrides operations to `dist`."""
  for slices, overrides in slice_overrides_seq:
    dist = _apply_single_step(dist, params_event_ndims, slices, overrides)
  return dist

def _apply_slice_sequence ( dist , params_event_ndims , slice_overrides_seq ) : for slices , overrides in slice_overrides_seq : dist = _apply_single_step ( dist , params_event_ndims , slices , overrides ) return dist

Applies a sequence of slice or copy - with - overrides operations to dist .


def batch_slice(dist, params_event_ndims, params_overrides, slices):
  """Slices `dist` along its batch dimensions. Helper for tfd.Distribution.

  Args:
    dist: A `tfd.Distribution` instance.
    params_event_ndims: A `dict` of `str->int` indicating the number of
      dimensions of a given parameter required to parameterize a single event.
    params_overrides: A `dict` of parameter overrides. (e.g. from
      `Distribution.copy`).
    slices: A `slice` or `int` or `int` `Tensor` or `tf.newaxis` or `tuple`
      thereof. (e.g. the argument of a `__getitem__` method).

  Returns:
    new_dist: A batch-sliced `tfd.Distribution`.
  """
  if not isinstance(slices, collections.Sequence):
    slices = (slices,)
  # We track the history of slice and copy(**param_overrides) in order to trace
  # back to the original distribution's source variables.
  orig_dist, slice_overrides_seq = getattr(dist, PROVENANCE_ATTR, (dist, []))
  slice_overrides_seq += [(slices, params_overrides)]
  # Re-doing the full sequence of slice+copy override work here enables
  # gradients all the way back to the original distribution's arguments.
  dist = _apply_slice_sequence(orig_dist, params_event_ndims,
                               slice_overrides_seq)
  setattr(dist, PROVENANCE_ATTR, (orig_dist, slice_overrides_seq))
  return dist

Slices dist along its batch dimensions . Helper for tfd . Distribution .


def fit(
    model_matrix,
    response,
    model,
    model_coefficients_start=None,
    predicted_linear_response_start=None,
    l2_regularizer=None,
    dispersion=None,
    offset=None,
    convergence_criteria_fn=None,
    learning_rate=None,
    fast_unsafe_numerics=True,
    maximum_iterations=None,
    name=None):
  """Runs multiple Fisher scoring steps.

  Args:
    model_matrix: (Batch of) `float`-like, matrix-shaped `Tensor` where each row
      represents a sample's features.
    response: (Batch of) vector-shaped `Tensor` where each element represents a
      sample's observed response (to the corresponding row of features). Must
      have same `dtype` as `model_matrix`.
    model: `tfp.glm.ExponentialFamily`-like instance which implicitly
      characterizes a negative log-likelihood loss by specifying the
      distribuion's `mean`, `gradient_mean`, and `variance`.
    model_coefficients_start: Optional (batch of) vector-shaped `Tensor`
      representing the initial model coefficients, one for each column in
      `model_matrix`. Must have same `dtype` as `model_matrix`.
      Default value: Zeros.
    predicted_linear_response_start: Optional `Tensor` with `shape`, `dtype`
      matching `response`; represents `offset` shifted initial linear
      predictions based on `model_coefficients_start`.
      Default value: `offset` if `model_coefficients is None`, and
      `tf.linalg.matvec(model_matrix, model_coefficients_start) + offset`
      otherwise.
    l2_regularizer: Optional scalar `Tensor` representing L2 regularization
      penalty, i.e.,
      `loss(w) = sum{-log p(y[i]|x[i],w) : i=1..n} + l2_regularizer ||w||_2^2`.
      Default value: `None` (i.e., no L2 regularization).
    dispersion: Optional (batch of) `Tensor` representing `response` dispersion,
      i.e., as in, `p(y|theta) := exp((y theta - A(theta)) / dispersion)`.
      Must broadcast with rows of `model_matrix`.
      Default value: `None` (i.e., "no dispersion").
    offset: Optional `Tensor` representing constant shift applied to
      `predicted_linear_response`.  Must broadcast to `response`.
      Default value: `None` (i.e., `tf.zeros_like(response)`).
    convergence_criteria_fn: Python `callable` taking:
      `is_converged_previous`, `iter_`, `model_coefficients_previous`,
      `predicted_linear_response_previous`, `model_coefficients_next`,
      `predicted_linear_response_next`, `response`, `model`, `dispersion` and
      returning a `bool` `Tensor` indicating that Fisher scoring has converged.
      See `convergence_criteria_small_relative_norm_weights_change` as an
      example function.
      Default value: `None` (i.e.,
      `convergence_criteria_small_relative_norm_weights_change`).
    learning_rate: Optional (batch of) scalar `Tensor` used to dampen iterative
      progress. Typically only needed if optimization diverges, should be no
      larger than `1` and typically very close to `1`.
      Default value: `None` (i.e., `1`).
    fast_unsafe_numerics: Optional Python `bool` indicating if faster, less
      numerically accurate methods can be employed for computing the weighted
      least-squares solution.
      Default value: `True` (i.e., "fast but possibly diminished accuracy").
    maximum_iterations: Optional maximum number of iterations of Fisher scoring
      to run; "and-ed" with result of `convergence_criteria_fn`.
      Default value: `None` (i.e., `infinity`).
    name: Python `str` used as name prefix to ops created by this function.
      Default value: `"fit"`.

  Returns:
    model_coefficients: (Batch of) vector-shaped `Tensor`; represents the
      fitted model coefficients, one for each column in `model_matrix`.
    predicted_linear_response: `response`-shaped `Tensor` representing linear
      predictions based on new `model_coefficients`, i.e.,
      `tf.linalg.matvec(model_matrix, model_coefficients) + offset`.
    is_converged: `bool` `Tensor` indicating that the returned
      `model_coefficients` met the `convergence_criteria_fn` criteria within the
      `maximum_iterations` limit.
    iter_: `int32` `Tensor` indicating the number of iterations taken.

  #### Example

  ```python
  from __future__ import print_function
  import numpy as np
  import tensorflow as tf
  import tensorflow_probability as tfp
  tfd = tfp.distributions

  def make_dataset(n, d, link, scale=1., dtype=np.float32):
    model_coefficients = tfd.Uniform(
        low=np.array(-1, dtype),
        high=np.array(1, dtype)).sample(d, seed=42)
    radius = np.sqrt(2.)
    model_coefficients *= radius / tf.linalg.norm(model_coefficients)
    model_matrix = tfd.Normal(
        loc=np.array(0, dtype),
        scale=np.array(1, dtype)).sample([n, d], seed=43)
    scale = tf.convert_to_tensor(scale, dtype)
    linear_response = tf.tensordot(
        model_matrix, model_coefficients, axes=[[1], [0]])
    if link == 'linear':
      response = tfd.Normal(loc=linear_response, scale=scale).sample(seed=44)
    elif link == 'probit':
      response = tf.cast(
          tfd.Normal(loc=linear_response, scale=scale).sample(seed=44) > 0,
          dtype)
    elif link == 'logit':
      response = tfd.Bernoulli(logits=linear_response).sample(seed=44)
    else:
      raise ValueError('unrecognized true link: {}'.format(link))
    return model_matrix, response, model_coefficients

  X, Y, w_true = make_dataset(n=int(1e6), d=100, link='probit')

  w, linear_response, is_converged, num_iter = tfp.glm.fit(
      model_matrix=X,
      response=Y,
      model=tfp.glm.BernoulliNormalCDF())
  log_likelihood = tfp.glm.BernoulliNormalCDF().log_prob(Y, linear_response)

  with tf.Session() as sess:
    [w_, linear_response_, is_converged_, num_iter_, Y_, w_true_,
     log_likelihood_] = sess.run([
        w, linear_response, is_converged, num_iter, Y, w_true,
        log_likelihood])

  print('is_converged: ', is_converged_)
  print('    num_iter: ', num_iter_)
  print('    accuracy: ', np.mean((linear_response_ > 0.) == Y_))
  print('    deviance: ', 2. * np.mean(log_likelihood_))
  print('||w0-w1||_2 / (1+||w0||_2): ', (np.linalg.norm(w_true_ - w_, ord=2) /
                                         (1. + np.linalg.norm(w_true_, ord=2))))

  # ==>
  # is_converged:  True
  #     num_iter:  6
  #     accuracy:  0.804382
  #     deviance:  -0.820746600628
  # ||w0-w1||_2 / (1+||w0||_2):  0.00619245105309
  ```

  """
  graph_deps = [model_matrix, response, model_coefficients_start,
                predicted_linear_response_start, dispersion, offset,
                learning_rate, maximum_iterations]
  with tf.compat.v1.name_scope(name, 'fit', graph_deps):
    [
        model_matrix,
        response,
        model_coefficients_start,
        predicted_linear_response_start,
        offset,
    ] = prepare_args(
        model_matrix,
        response,
        model_coefficients_start,
        predicted_linear_response_start,
        offset)
    if convergence_criteria_fn is None:
      convergence_criteria_fn = (
          convergence_criteria_small_relative_norm_weights_change())

    def _body(
        is_converged_previous,
        iter_,
        model_coefficients_previous,
        predicted_linear_response_previous):
      """`tf.while_loop` body."""
      model_coefficients_next, predicted_linear_response_next = fit_one_step(
          model_matrix,
          response,
          model,
          model_coefficients_previous,
          predicted_linear_response_previous,
          l2_regularizer,
          dispersion,
          offset,
          learning_rate,
          fast_unsafe_numerics)
      is_converged_next = convergence_criteria_fn(
          is_converged_previous=is_converged_previous,
          iter_=iter_,
          model_coefficients_previous=model_coefficients_previous,
          predicted_linear_response_previous=predicted_linear_response_previous,
          model_coefficients_next=model_coefficients_next,
          predicted_linear_response_next=predicted_linear_response_next,
          response=response,
          model=model,
          dispersion=dispersion)
      return [
          is_converged_next,
          iter_ + 1,
          model_coefficients_next,
          predicted_linear_response_next,
      ]

    # while not converged:
    #   fit_one_step
    [
        is_converged,
        iter_,
        model_coefficients,
        predicted_linear_response,
    ] = tf.while_loop(
        cond=lambda is_converged, *args: tf.logical_not(is_converged),
        body=_body,
        loop_vars=[
            tf.zeros([], np.bool),   # is_converged
            tf.zeros([], np.int32),  # iter_
            model_coefficients_start,
            predicted_linear_response_start,
        ],
        maximum_iterations=maximum_iterations)

    return [
        model_coefficients,
        predicted_linear_response,
        is_converged,
        iter_
    ]


def convergence_criteria_small_relative_norm_weights_change(
    tolerance=1e-5,
    norm_order=2):
  """Returns Python `callable` which indicates fitting procedure has converged.

  Writing old, new `model_coefficients` as `w0`, `w1`, this function
  defines convergence as,

  ```python
  relative_euclidean_norm = (tf.norm(w0 - w1, ord=2, axis=-1) /
                             (1. + tf.norm(w0, ord=2, axis=-1)))
  reduce_all(relative_euclidean_norm < tolerance)
  ```

  where `tf.norm(x, ord=2)` denotes the [Euclidean norm](
  https://en.wikipedia.org/wiki/Norm_(mathematics)#Euclidean_norm) of `x`.

  Args:
    tolerance: `float`-like `Tensor` indicating convergence, i.e., when
      max relative Euclidean norm weights difference < tolerance`.
      Default value: `1e-5`.
    norm_order: Order of the norm. Default value: `2` (i.e., "Euclidean norm".)

  Returns:
    convergence_criteria_fn: Python `callable` which returns `bool` `Tensor`
      indicated fitting procedure has converged. (See inner function
      specification for argument signature.)
      Default value: `1e-5`.
  """
  def convergence_criteria_fn(
      is_converged_previous,  # pylint: disable=unused-argument
      iter_,
      model_coefficients_previous,
      predicted_linear_response_previous,  # pylint: disable=unused-argument
      model_coefficients_next,
      predicted_linear_response_next,  # pylint: disable=unused-argument
      response,  # pylint: disable=unused-argument
      model,  # pylint: disable=unused-argument
      dispersion):  # pylint: disable=unused-argument
    """Returns `bool` `Tensor` indicating if fitting procedure has converged.

    Args:
      is_converged_previous: "old" convergence results.
      iter_: Iteration number.
      model_coefficients_previous: "old" `model_coefficients`.
      predicted_linear_response_previous: "old" `predicted_linear_response`.
      model_coefficients_next: "new" `model_coefficients`.
      predicted_linear_response_next: "new: `predicted_linear_response`.
      response: (Batch of) vector-shaped `Tensor` where each element represents
        a sample's observed response (to the corresponding row of features).
        Must have same `dtype` as `model_matrix`.
      model: `tfp.glm.ExponentialFamily`-like instance used to construct the
        negative log-likelihood loss, gradient, and expected Hessian (i.e., the
        Fisher information matrix).
      dispersion: `Tensor` representing `response` dispersion, i.e., as in:
        `p(y|theta) := exp((y theta - A(theta)) / dispersion)`. Must broadcast
        with rows of `model_matrix`.
        Default value: `None` (i.e., "no dispersion").

    Returns:
      is_converged: `bool` `Tensor`.
    """
    relative_euclidean_norm = (
        tf.norm(
            tensor=model_coefficients_previous - model_coefficients_next,
            ord=norm_order,
            axis=-1) /
        (1. +
         tf.norm(tensor=model_coefficients_previous, ord=norm_order, axis=-1)))
    return (iter_ > 0) & tf.reduce_all(
        input_tensor=relative_euclidean_norm < tolerance)

  return convergence_criteria_fn

Returns Python callable which indicates fitting procedure has converged .


def prepare_args(model_matrix,
                 response,
                 model_coefficients,
                 predicted_linear_response,
                 offset,
                 name=None):
  """Helper to `fit` which sanitizes input args.

  Args:
    model_matrix: (Batch of) `float`-like, matrix-shaped `Tensor` where each row
      represents a sample's features.
    response: (Batch of) vector-shaped `Tensor` where each element represents a
      sample's observed response (to the corresponding row of features). Must
      have same `dtype` as `model_matrix`.
    model_coefficients: Optional (batch of) vector-shaped `Tensor` representing
      the model coefficients, one for each column in `model_matrix`. Must have
      same `dtype` as `model_matrix`.
      Default value: `tf.zeros(tf.shape(model_matrix)[-1], model_matrix.dtype)`.
    predicted_linear_response: Optional `Tensor` with `shape`, `dtype` matching
      `response`; represents `offset` shifted initial linear predictions based
      on current `model_coefficients`.
      Default value: `offset` if `model_coefficients is None`, and
      `tf.linalg.matvec(model_matrix, model_coefficients_start) + offset`
      otherwise.
    offset: Optional `Tensor` with `shape`, `dtype` matching `response`;
      represents constant shift applied to `predicted_linear_response`.
      Default value: `None` (i.e., `tf.zeros_like(response)`).
    name: Python `str` used as name prefix to ops created by this function.
      Default value: `"prepare_args"`.

  Returns:
    model_matrix: A `Tensor` with `shape`, `dtype` and values of the
      `model_matrix` argument.
    response: A `Tensor` with `shape`, `dtype` and values of the
      `response` argument.
    model_coefficients_start: A `Tensor` with `shape`, `dtype` and
      values of the `model_coefficients_start` argument if specified.
      A (batch of) vector-shaped `Tensors` with `dtype` matching `model_matrix`
      containing the default starting point otherwise.
    predicted_linear_response:  A `Tensor` with `shape`, `dtype` and
      values of the `predicted_linear_response` argument if specified.
      A `Tensor` with `shape`, `dtype` matching `response` containing the
      default value otherwise.
    offset: A `Tensor` with `shape`, `dtype` and values of the `offset` argument
      if specified or `None` otherwise.
  """
  graph_deps = [model_matrix, response, model_coefficients,
                predicted_linear_response, offset]
  with tf.compat.v1.name_scope(name, 'prepare_args', graph_deps):
    dtype = dtype_util.common_dtype(graph_deps, np.float32)

    model_matrix = tf.convert_to_tensor(
        value=model_matrix, dtype=dtype, name='model_matrix')

    if offset is not None:
      offset = tf.convert_to_tensor(value=offset, dtype=dtype, name='offset')

    response = tf.convert_to_tensor(
        value=response, dtype=dtype, name='response')

    use_default_model_coefficients = model_coefficients is None
    if use_default_model_coefficients:
      # User did not supply model coefficients; assume they're all zero.
      batch_shape = tf.shape(input=model_matrix)[:-2]
      num_columns = tf.shape(input=model_matrix)[-1]
      model_coefficients = tf.zeros(
          shape=tf.concat([batch_shape, [num_columns]], axis=0),
          dtype=dtype, name='model_coefficients')
    else:
      # User did supply model coefficients; convert to Tensor in case it's
      # numpy or literal.
      model_coefficients = tf.convert_to_tensor(
          value=model_coefficients, dtype=dtype, name='model_coefficients')

    if predicted_linear_response is None:
      if use_default_model_coefficients:
        # Since we're using zeros for model_coefficients, we know the predicted
        # linear response will also be all zeros.
        if offset is None:
          predicted_linear_response = tf.zeros_like(
              response, dtype, name='predicted_linear_response')
        else:
          predicted_linear_response = tf.broadcast_to(
              offset,
              tf.shape(input=response),
              name='predicted_linear_response')
      else:
        # We were given model_coefficients but not the predicted linear
        # response.
        predicted_linear_response = calculate_linear_predictor(
            model_matrix, model_coefficients, offset)
    else:
      predicted_linear_response = tf.convert_to_tensor(
          value=predicted_linear_response,
          dtype=dtype,
          name='predicted_linear_response')

  return [
      model_matrix,
      response,
      model_coefficients,
      predicted_linear_response,
      offset,
  ]

Helper to fit which sanitizes input args .


def num_cols(x):
  """Returns number of cols in a given `Tensor`."""
  if tf.compat.dimension_value(x.shape[-1]) is not None:
    return tf.compat.dimension_value(x.shape[-1])
  return tf.shape(input=x)[-1]

def num_cols ( x ) : if tf . compat . dimension_value ( x . shape [ - 1 ] ) is not None : return tf . compat . dimension_value ( x . shape [ - 1 ] ) return tf . shape ( input = x ) [ - 1 ]

Returns number of cols in a given Tensor .


def _prefer_static(original_fn, static_fn):
  """Wraps original_fn, preferring to call static_fn when inputs are static."""
  original_spec = tf_inspect.getfullargspec(original_fn)
  static_spec = tf_inspect.getfullargspec(static_fn)
  if original_spec != static_spec:
    raise ValueError(
        'Arg specs do not match: original={}, static={}, fn={}'.format(
            original_spec, static_spec, original_fn))
  @decorator.decorator
  def wrap(wrapped_fn, *args, **kwargs):
    del wrapped_fn
    [args_, kwargs_], all_static = _maybe_get_static_args([args, kwargs])
    if all_static:
      return static_fn(*args_, **kwargs_)
    return original_fn(*args, **kwargs)
  return wrap(original_fn)

Wraps original_fn preferring to call static_fn when inputs are static .


def _copy_docstring(original_fn, new_fn):
  """Wraps new_fn with the doc of original_fn."""
  original_spec = tf_inspect.getfullargspec(original_fn)
  new_spec = tf_inspect.getfullargspec(new_fn)
  if original_spec != new_spec:
    raise ValueError(
        'Arg specs do not match: original={}, new={}, fn={}'.format(
            original_spec, new_spec, original_fn))
  @decorator.decorator
  def wrap(wrapped_fn, *args, **kwargs):
    del wrapped_fn
    return new_fn(*args, **kwargs)
  return wrap(original_fn)

Wraps new_fn with the doc of original_fn .


def _get_static_predicate(pred):
  """Helper function for statically evaluating predicates in `cond`."""
  if pred in {0, 1}:  # Accept 1/0 as valid boolean values
    pred_value = bool(pred)
  elif isinstance(pred, bool):
    pred_value = pred
  elif isinstance(pred, tf.Tensor):
    pred_value = tf.get_static_value(pred)

    # TODO(jamieas): remove the dependency on `pywrap_tensorflow`.
    # pylint: disable=protected-access
    if pred_value is None:
      pred_value = c_api.TF_TryEvaluateConstant_wrapper(pred.graph._c_graph,
                                                        pred._as_tf_output())
    # pylint: enable=protected-access

  else:
    raise TypeError('`pred` must be a Tensor, or a Python bool, or 1 or 0. '
                    'Found instead: {}'.format(pred))
  return pred_value

Helper function for statically evaluating predicates in cond .


def rank_from_shape(shape_tensor_fn, tensorshape=None):
  """Computes `rank` given a `Tensor`'s `shape`."""

  if tensorshape is None:
    shape_tensor = (shape_tensor_fn() if callable(shape_tensor_fn)
                    else shape_tensor_fn)
    if (hasattr(shape_tensor, 'shape') and
        hasattr(shape_tensor.shape, 'num_elements')):
      ndims_ = tensorshape_util.num_elements(shape_tensor.shape)
    else:
      ndims_ = len(shape_tensor)
    ndims_fn = lambda: tf.size(input=shape_tensor)
  else:
    ndims_ = tensorshape_util.rank(tensorshape)
    ndims_fn = lambda: tf.size(input=shape_tensor_fn()  # pylint: disable=g-long-lambda
                               if callable(shape_tensor_fn)
                               else shape_tensor_fn)
  return ndims_fn() if ndims_ is None else ndims_

def
rank_from_shape
(
shape_tensor_fn
,
tensorshape
=
)
:
if
tensorshape
is
:
shape_tensor
=
(
shape_tensor_fn
(
)
if
callable
(
shape_tensor_fn
)
else
shape_tensor_fn
)
if
(
hasattr
(
shape_tensor
,
'shape'
)
and
hasattr
(
shape_tensor
.
shape
,
'num_elements'
)
)
:
ndims_
=
tensorshape_util
.
num_elements
(
shape_tensor
.
shape
)
else
:
ndims_
=
len
(
shape_tensor
)
ndims_fn
=
lambda
:
tf
.
size
(
input
=
shape_tensor
)
else
:
ndims_
=
tensorshape_util
.
rank
(
tensorshape
)
ndims_fn
=
lambda
:
tf
.
size
(
input
=
shape_tensor_fn
(
)
# pylint: disable=g-long-lambda
if
callable
(
shape_tensor_fn
)
else
shape_tensor_fn
)
return
ndims_fn
(
)
if
ndims_
is
else
ndims_


def case(pred_fn_pairs, default=None, exclusive=False, name='smart_case'):
  """Like tf.case, except attempts to statically evaluate predicates.

  If any predicate in `pred_fn_pairs` is a bool or has a constant value, the
  associated callable will be called or omitted depending on its value.
  Otherwise this functions like tf.case.

  Args:
    pred_fn_pairs: Dict or list of pairs of a boolean scalar tensor and a
                   callable which returns a list of tensors.
    default: Optional callable that returns a list of tensors.
    exclusive: True iff at most one predicate is allowed to evaluate to `True`.
    name: A name for this operation (optional).

  Returns:
    The tensors returned by the first pair whose predicate evaluated to True, or
    those returned by `default` if none does.

  Raises:
    TypeError: If `pred_fn_pairs` is not a list/dictionary.
    TypeError: If `pred_fn_pairs` is a list but does not contain 2-tuples.
    TypeError: If `fns[i]` is not callable for any i, or `default` is not
               callable.
  """
  return control_flow_ops._case_helper(  # pylint: disable=protected-access
      cond, pred_fn_pairs, default, exclusive, name, allow_python_preds=True)

def
case
(
pred_fn_pairs
,
default
=
,
exclusive
=
False
,
name
=
'smart_case'
)
:
return
control_flow_ops
.
_case_helper
(
# pylint: disable=protected-access
cond
,
pred_fn_pairs
,
default
,
exclusive
,
name
,
allow_python_preds
=
True
)

Like tf . case except attempts to statically evaluate predicates .


def _name_scope(self, name=None, default_name=None, values=None):
    """Helper function to standardize op scope."""
    with tf.compat.v1.name_scope(self.name):
      with tf.compat.v1.name_scope(
          name, default_name, values=values or []) as scope:
        yield scope

def _name_scope ( self , name = None , default_name = None , values = None ) : with tf . compat . v1 . name_scope ( self . name ) : with tf . compat . v1 . name_scope ( name , default_name , values = values or [ ] ) as scope : yield scope

Helper function to standardize op scope .