def forecast(model,
             observed_time_series,
             parameter_samples,
             num_steps_forecast):
  """Construct predictive distribution over future observations.

  Given samples from the posterior over parameters, return the predictive
  distribution over future observations for num_steps_forecast timesteps.

  Args:
    model: An instance of `StructuralTimeSeries` representing a
      time-series model. This represents a joint distribution over
      time-series and their parameters with batch shape `[b1, ..., bN]`.
    observed_time_series: `float` `Tensor` of shape
      `concat([sample_shape, model.batch_shape, [num_timesteps, 1]])` where
      `sample_shape` corresponds to i.i.d. observations, and the trailing `[1]`
      dimension may (optionally) be omitted if `num_timesteps > 1`. May
      optionally be an instance of `tfp.sts.MaskedTimeSeries` including a
      mask `Tensor` to encode the locations of missing observations.
    parameter_samples: Python `list` of `Tensors` representing posterior samples
      of model parameters, with shapes `[concat([[num_posterior_draws],
      param.prior.batch_shape, param.prior.event_shape]) for param in
      model.parameters]`. This may optionally also be a map (Python `dict`) of
      parameter names to `Tensor` values.
    num_steps_forecast: scalar `int` `Tensor` number of steps to forecast.

  Returns:
    forecast_dist: a `tfd.MixtureSameFamily` instance with event shape
      [num_steps_forecast, 1] and batch shape
      `concat([sample_shape, model.batch_shape])`, with `num_posterior_draws`
      mixture components.

  #### Examples

  Suppose we've built a model and fit it to data using HMC:

  ```python
    day_of_week = tfp.sts.Seasonal(
        num_seasons=7,
        observed_time_series=observed_time_series,
        name='day_of_week')
    local_linear_trend = tfp.sts.LocalLinearTrend(
        observed_time_series=observed_time_series,
        name='local_linear_trend')
    model = tfp.sts.Sum(components=[day_of_week, local_linear_trend],
                        observed_time_series=observed_time_series)

    samples, kernel_results = tfp.sts.fit_with_hmc(model, observed_time_series)
  ```

  Passing the posterior samples into `forecast`, we construct a forecast
  distribution:

  ```python
    forecast_dist = tfp.sts.forecast(model, observed_time_series,
                                     parameter_samples=samples,
                                     num_steps_forecast=50)

    forecast_mean = forecast_dist.mean()[..., 0]  # shape: [50]
    forecast_scale = forecast_dist.stddev()[..., 0]  # shape: [50]
    forecast_samples = forecast_dist.sample(10)[..., 0]  # shape: [10, 50]
  ```

  If using variational inference instead of HMC, we'd construct a forecast using
  samples from the variational posterior:

  ```python
    (variational_loss,
     variational_distributions) = tfp.sts.build_factored_variational_loss(
       model=model, observed_time_series=observed_time_series)

    # OMITTED: take steps to optimize variational loss

    samples = {k: q.sample(30) for (k, q) in variational_distributions.items()}
    forecast_dist = tfp.sts.forecast(model, observed_time_series,
                                         parameter_samples=samples,
                                         num_steps_forecast=50)
  ```

  We can visualize the forecast by plotting:

  ```python
    from matplotlib import pylab as plt
    def plot_forecast(observed_time_series,
                      forecast_mean,
                      forecast_scale,
                      forecast_samples):
      plt.figure(figsize=(12, 6))

      num_steps = observed_time_series.shape[-1]
      num_steps_forecast = forecast_mean.shape[-1]
      num_steps_train = num_steps - num_steps_forecast

      c1, c2 = (0.12, 0.47, 0.71), (1.0, 0.5, 0.05)
      plt.plot(np.arange(num_steps), observed_time_series,
               lw=2, color=c1, label='ground truth')

      forecast_steps = np.arange(num_steps_train,
                       num_steps_train+num_steps_forecast)
      plt.plot(forecast_steps, forecast_samples.T, lw=1, color=c2, alpha=0.1)
      plt.plot(forecast_steps, forecast_mean, lw=2, ls='--', color=c2,
               label='forecast')
      plt.fill_between(forecast_steps,
                       forecast_mean - 2 * forecast_scale,
                       forecast_mean + 2 * forecast_scale, color=c2, alpha=0.2)

      plt.xlim([0, num_steps])
      plt.legend()

    plot_forecast(observed_time_series,
                  forecast_mean=forecast_mean,
                  forecast_scale=forecast_scale,
                  forecast_samples=forecast_samples)
  ```

  """

  with tf.compat.v1.name_scope(
      'forecast',
      values=[observed_time_series, parameter_samples, num_steps_forecast]):
    [
        observed_time_series,
        mask
    ] = sts_util.canonicalize_observed_time_series_with_mask(
        observed_time_series)

    # Run filtering over the observed timesteps to extract the
    # latent state posterior at timestep T+1 (i.e., the final
    # filtering distribution, pushed through the transition model).
    # This is the prior for the forecast model ("today's prior
    # is yesterday's posterior").
    num_observed_steps = dist_util.prefer_static_value(
        tf.shape(input=observed_time_series))[-2]
    observed_data_ssm = model.make_state_space_model(
        num_timesteps=num_observed_steps, param_vals=parameter_samples)
    (_, _, _, predictive_means, predictive_covs, _, _
    ) = observed_data_ssm.forward_filter(observed_time_series, mask=mask)

    # Build a batch of state-space models over the forecast period. Because
    # we'll use MixtureSameFamily to mix over the posterior draws, we need to
    # do some shenanigans to move the `[num_posterior_draws]` batch dimension
    # from the leftmost to the rightmost side of the model's batch shape.
    # TODO(b/120245392): enhance `MixtureSameFamily` to reduce along an
    # arbitrary axis, and eliminate `move_dimension` calls here.
    parameter_samples = model._canonicalize_param_vals_as_map(parameter_samples)  # pylint: disable=protected-access
    parameter_samples_with_reordered_batch_dimension = {
        param.name: dist_util.move_dimension(
            parameter_samples[param.name],
            0, -(1 + _prefer_static_event_ndims(param.prior)))
        for param in model.parameters}
    forecast_prior = tfd.MultivariateNormalFullCovariance(
        loc=dist_util.move_dimension(predictive_means[..., -1, :], 0, -2),
        covariance_matrix=dist_util.move_dimension(
            predictive_covs[..., -1, :, :], 0, -3))

    # Ugly hack: because we moved `num_posterior_draws` to the trailing (rather
    # than leading) dimension of parameters, the parameter batch shapes no
    # longer broadcast against the `constant_offset` attribute used in `sts.Sum`
    # models. We fix this by manually adding an extra broadcasting dim to
    # `constant_offset` if present.
    # The root cause of this hack is that we mucked with param dimensions above
    # and are now passing params that are 'invalid' in the sense that they don't
    # match the shapes of the model's param priors. The fix (as above) will be
    # to update MixtureSameFamily so we can avoid changing param dimensions
    # altogether.
    # TODO(b/120245392): enhance `MixtureSameFamily` to reduce along an
    # arbitrary axis, and eliminate this hack.
    kwargs = {}
    if hasattr(model, 'constant_offset'):
      kwargs['constant_offset'] = tf.convert_to_tensor(
          value=model.constant_offset,
          dtype=forecast_prior.dtype)[..., tf.newaxis]

    # We assume that any STS model that has a `constant_offset` attribute
    # will allow it to be overridden as a kwarg. This is currently just
    # `sts.Sum`.
    # TODO(b/120245392): when kwargs hack is removed, switch back to calling
    # the public version of `_make_state_space_model`.
    forecast_ssm = model._make_state_space_model(  # pylint: disable=protected-access
        num_timesteps=num_steps_forecast,
        param_map=parameter_samples_with_reordered_batch_dimension,
        initial_state_prior=forecast_prior,
        initial_step=num_observed_steps,
        **kwargs)

    num_posterior_draws = dist_util.prefer_static_value(
        forecast_ssm.batch_shape_tensor())[-1]
    return tfd.MixtureSameFamily(
        mixture_distribution=tfd.Categorical(
            logits=tf.zeros([num_posterior_draws], dtype=forecast_ssm.dtype)),
        components_distribution=forecast_ssm)

Construct predictive distribution over future observations .


def _max_mask_non_finite(x, axis=-1, keepdims=False, mask=0):
  """Returns `max` or `mask` if `max` is not finite."""
  m = np.max(x, axis=_astuple(axis), keepdims=keepdims)
  needs_masking = ~np.isfinite(m)
  if needs_masking.ndim > 0:
    m[needs_masking] = mask
  elif needs_masking:
    m = mask
  return m

def _max_mask_non_finite ( x , axis = - 1 , keepdims = False , mask = 0 ) : m = np . max ( x , axis = _astuple ( axis ) , keepdims = keepdims ) needs_masking = ~ np . isfinite ( m ) if needs_masking . ndim > 0 : m [ needs_masking ] = mask elif needs_masking : m = mask return m

Returns max or mask if max is not finite .


def assert_finite(x, data=None, summarize=None, message=None, name=None):
  """Assert all elements of `x` are finite.

  Args:
    x:  Numeric `Tensor`.
    data:  The tensors to print out if the condition is False.  Defaults to
      error message and first few entries of `x`.
    summarize: Print this many entries of each tensor.
    message: A string to prefix to the default message.
    name: A name for this operation (optional).
      Defaults to "assert_finite".

  Returns:
    Op raising `InvalidArgumentError` unless `x` has specified rank or lower.
    If static checks determine `x` has correct rank, a `no_op` is returned.

  Raises:
    ValueError:  If static checks determine `x` has wrong rank.
  """
  with tf.compat.v2.name_scope(name or 'assert_finite'):
    x_ = tf.get_static_value(x)
    if x_ is not None:
      if ~np.all(np.isfinite(x_)):
        raise ValueError(message)
      return x
    assertion = tf.compat.v1.assert_equal(
        tf.math.is_finite(x), tf.ones_like(x, tf.bool),
        data=data, summarize=summarize, message=message)
    with tf.control_dependencies([assertion]):
      return tf.identity(x)

def assert_finite ( x , data = None , summarize = None , message = None , name = None ) : with tf . compat . v2 . name_scope ( name or 'assert_finite' ) : x_ = tf . get_static_value ( x ) if x_ is not None : if ~ np . all ( np . isfinite ( x_ ) ) : raise ValueError ( message ) return x assertion = tf . compat . v1 . assert_equal ( tf . math . is_finite ( x ) , tf . ones_like ( x , tf . bool ) , data = data , summarize = summarize , message = message ) with tf . control_dependencies ( [ assertion ] ) : return tf . identity ( x )


def assert_rank_at_most(x, rank, data=None, summarize=None, message=None,
                        name=None):
  """Assert `x` has rank equal to `rank` or smaller.

  Example of adding a dependency to an operation:

  ```python
  with tf.control_dependencies([tf.assert_rank_at_most(x, 2)]):
    output = tf.reduce_sum(x)
  ```

  Args:
    x:  Numeric `Tensor`.
    rank:  Scalar `Tensor`.
    data:  The tensors to print out if the condition is False.  Defaults to
      error message and first few entries of `x`.
    summarize: Print this many entries of each tensor.
    message: A string to prefix to the default message.
    name: A name for this operation (optional).
      Defaults to "assert_rank_at_most".

  Returns:
    Op raising `InvalidArgumentError` unless `x` has specified rank or lower.
    If static checks determine `x` has correct rank, a `no_op` is returned.

  Raises:
    ValueError:  If static checks determine `x` has wrong rank.
  """
  with tf.compat.v2.name_scope(name or 'assert_rank_at_most'):
    return tf.compat.v1.assert_less_equal(
        tf.rank(x), rank, data=data, summarize=summarize, message=message)

def assert_rank_at_most ( x , rank , data = None , summarize = None , message = None , name = None ) : with tf . compat . v2 . name_scope ( name or 'assert_rank_at_most' ) : return tf . compat . v1 . assert_less_equal ( tf . rank ( x ) , rank , data = data , summarize = summarize , message = message )

Assert x has rank equal to rank or smaller .


def _event_size(event_shape, name=None):
  """Computes the number of elements in a tensor with shape `event_shape`.

  Args:
    event_shape: A tensor shape.
    name: The name to use for the tensor op to compute the number of elements
      (if such an op needs to be created).

  Returns:
    event_size: The number of elements in `tensor_shape`.  Returns a numpy int
    when the number of elements can be computed immediately.  Otherwise, returns
    a scalar tensor.
  """
  with tf.compat.v1.name_scope(name, 'event_size', [event_shape]):
    event_shape = tf.convert_to_tensor(
        value=event_shape, dtype=tf.int32, name='event_shape')

    event_shape_const = tf.get_static_value(event_shape)
    if event_shape_const is not None:
      return np.prod(event_shape_const)
    else:
      return tf.reduce_prod(input_tensor=event_shape)

def _event_size ( event_shape , name = None ) : with tf . compat . v1 . name_scope ( name , 'event_size' , [ event_shape ] ) : event_shape = tf . convert_to_tensor ( value = event_shape , dtype = tf . int32 , name = 'event_shape' ) event_shape_const = tf . get_static_value ( event_shape ) if event_shape_const is not None : return np . prod ( event_shape_const ) else : return tf . reduce_prod ( input_tensor = event_shape )

Computes the number of elements in a tensor with shape event_shape .


def _eval_all_one_hot(fn, dist, name=None):
  """OneHotCategorical helper computing probs, cdf, etc over its support."""
  with tf.compat.v1.name_scope(name, 'eval_all_one_hot'):
    event_size = dist.event_shape_tensor()[-1]
    batch_ndims = tf.size(input=dist.batch_shape_tensor())
    # Reshape `eye(d)` to: `[d] + [1]*batch_ndims + [d]`.
    x = tf.reshape(
        tf.eye(event_size, dtype=dist.dtype),
        shape=tf.pad(
            tensor=tf.ones(batch_ndims, tf.int32),
            paddings=[[1, 1]],
            constant_values=event_size))
    # Compute `fn(x)` then cyclically left-transpose one dim.
    perm = tf.pad(tensor=tf.range(1, batch_ndims + 1), paddings=[[0, 1]])
    return tf.transpose(a=fn(dist, x), perm=perm)

def _eval_all_one_hot ( fn , dist , name = None ) : with tf . compat . v1 . name_scope ( name , 'eval_all_one_hot' ) : event_size = dist . event_shape_tensor ( ) [ - 1 ] batch_ndims = tf . size ( input = dist . batch_shape_tensor ( ) ) x = tf . reshape ( tf . eye ( event_size , dtype = dist . dtype ) , shape = tf . pad ( tensor = tf . ones ( batch_ndims , tf . int32 ) , paddings = [ [ 1 , 1 ] ] , constant_values = event_size ) ) perm = tf . pad ( tensor = tf . range ( 1 , batch_ndims + 1 ) , paddings = [ [ 0 , 1 ] ] ) return tf . transpose ( a = fn ( dist , x ) , perm = perm )

OneHotCategorical helper computing probs cdf etc over its support .


def _get_convert_to_tensor_fn(identifier):
  """Return a convert-to-tensor func, given a name, config, callable, etc."""
  if identifier is None:
    return None

  if isinstance(identifier, six.string_types):
    identifier = str(identifier)
    return _deserialize(identifier)

  if isinstance(identifier, dict):
    return _deserialize(identifier)

  if isinstance(identifier, property):
    identifier = identifier.fget
  if callable(identifier):
    return identifier

  raise ValueError('Could not interpret '
                   'convert-to-tensor function identifier:', identifier)

Return a convert - to - tensor func given a name config callable etc .


def params_size(num_components, component_params_size, name=None):
    """Number of `params` needed to create a `MixtureSameFamily` distribution.

    Arguments:
      num_components: Number of component distributions in the mixture
        distribution.
      component_params_size: Number of parameters needed to create a single
        component distribution.
      name: The name to use for the op to compute the number of parameters
        (if such an op needs to be created).

    Returns:
     params_size: The number of parameters needed to create the mixture
       distribution.
    """
    with tf.compat.v1.name_scope(name, 'MixtureSameFamily_params_size',
                                 [num_components, component_params_size]):
      num_components = tf.convert_to_tensor(
          value=num_components, name='num_components', dtype_hint=tf.int32)
      component_params_size = tf.convert_to_tensor(
          value=component_params_size, name='component_params_size')

      num_components = dist_util.prefer_static_value(num_components)
      component_params_size = dist_util.prefer_static_value(
          component_params_size)

      return num_components + num_components * component_params_size

def params_size ( num_components , component_params_size , name = None ) : with tf . compat . v1 . name_scope ( name , 'MixtureSameFamily_params_size' , [ num_components , component_params_size ] ) : num_components = tf . convert_to_tensor ( value = num_components , name = 'num_components' , dtype_hint = tf . int32 ) component_params_size = tf . convert_to_tensor ( value = component_params_size , name = 'component_params_size' ) num_components = dist_util . prefer_static_value ( num_components ) component_params_size = dist_util . prefer_static_value ( component_params_size ) return num_components + num_components * component_params_size

Number of params needed to create a MixtureSameFamily distribution .


def get_next_interceptor():
  """Yields the top-most interceptor on the thread-local interceptor stack.

  Operations may be intercepted by multiple nested interceptors. Once reached,
  an operation can be forwarded through nested interceptors until resolved.
  To allow for nesting, implement interceptors by re-wrapping their first
  argument (`f`) as an `interceptable`. To avoid nesting, manipulate the
  computation without using `interceptable`.

  This function allows for nesting by manipulating the thread-local interceptor
  stack, so that operations are intercepted in the order of interceptor nesting.

  #### Examples

  ```python
  from tensorflow_probability import edward2 as ed

  def model():
    x = ed.Normal(loc=0., scale=1., name="x")
    y = ed.Normal(loc=x, scale=1., name="y")
    return x + y

  def double(f, *args, **kwargs):
    return 2. * interceptable(f)(*args, **kwargs)

  def set_y(f, *args, **kwargs):
    if kwargs.get("name") == "y":
      kwargs["value"] = 0.42
    return interceptable(f)(*args, **kwargs)

  with interception(double):
    with interception(set_y):
      z = model()
  ```

  This will firstly put `double` on the stack, and then `set_y`,
  resulting in the stack:
  (TOP) set_y -> double -> apply (BOTTOM)

  The execution of `model` is then (top lines are current stack state):
  1) (TOP) set_y -> double -> apply (BOTTOM);
  `ed.Normal(0., 1., "x")` is intercepted by `set_y`, and as the name is not "y"
  the operation is simply forwarded to the next interceptor on the stack.

  2) (TOP) double -> apply (BOTTOM);
  `ed.Normal(0., 1., "x")` is intercepted by `double`, to produce
  `2*ed.Normal(0., 1., "x")`, with the operation being forwarded down the stack.

  3) (TOP) apply (BOTTOM);
  `ed.Normal(0., 1., "x")` is intercepted by `apply`, which simply calls the
  constructor.

  (At this point, the nested calls to `get_next_interceptor()`, produced by
  forwarding operations, exit, and the current stack is again:
  (TOP) set_y -> double -> apply (BOTTOM))

  4) (TOP) set_y -> double -> apply (BOTTOM);
  `ed.Normal(0., 1., "y")` is intercepted by `set_y`,
  the value of `y` is set to 0.42 and the operation is forwarded down the stack.

  5) (TOP) double -> apply (BOTTOM);
  `ed.Normal(0., 1., "y")` is intercepted by `double`, to produce
  `2*ed.Normal(0., 1., "y")`, with the operation being forwarded down the stack.

  6) (TOP) apply (BOTTOM);
  `ed.Normal(0., 1., "y")` is intercepted by `apply`, which simply calls the
  constructor.

  The final values for `x` and `y` inside of `model()` are tensors where `x` is
  a random draw from Normal(0., 1.) doubled, and `y` is a constant 0.84, thus
  z = 2 * Normal(0., 1.) + 0.84.
  """
  try:
    interceptor = _interceptor_stack.stack.pop()
    yield interceptor
  finally:
    _interceptor_stack.stack.append(interceptor)

def get_next_interceptor ( ) : try : interceptor = _interceptor_stack . stack . pop ( ) yield interceptor finally : _interceptor_stack . stack . append ( interceptor )

Yields the top - most interceptor on the thread - local interceptor stack .


def interceptable(func):
  """Decorator that wraps `func` so that its execution is intercepted.

  The wrapper passes `func` to the interceptor for the current thread.

  If there is no next interceptor, we perform an "immediate" call to `func`.
  That is, `func` terminates without forwarding its execution to another
  interceptor.

  Args:
    func: Function to wrap.

  Returns:
    The decorated function.
  """
  @functools.wraps(func)
  def func_wrapped(*args, **kwargs):
    with get_next_interceptor() as interceptor:
      return interceptor(func, *args, **kwargs)

  return func_wrapped

def interceptable ( func ) : @ functools . wraps ( func ) def func_wrapped ( * args , ** kwargs ) : with get_next_interceptor ( ) as interceptor : return interceptor ( func , * args , ** kwargs ) return func_wrapped

Decorator that wraps func so that its execution is intercepted .


def tape():
  """Context manager for recording interceptable executions onto a tape.

  Similar to `tf.GradientTape`, operations are recorded if they are executed
  within this context manager. In addition, the operation must be registered
  (wrapped) as `ed.interceptable`.

  Yields:
    tape: OrderedDict where operations are recorded in sequence. Keys are
      the `name` keyword argument to the operation (typically, a random
      variable's `name`) and values are the corresponding output of the
      operation. If the operation has no name, it is not recorded.

  #### Examples

  ```python
  from tensorflow_probability import edward2 as ed

  def probabilistic_matrix_factorization():
    users = ed.Normal(0., 1., sample_shape=[5000, 128], name="users")
    items = ed.Normal(0., 1., sample_shape=[7500, 128], name="items")
    ratings = ed.Normal(loc=tf.matmul(users, items, transpose_b=True),
                        scale=0.1,
                        name="ratings")
    return ratings

  with ed.tape() as model_tape:
    ratings = probabilistic_matrix_factorization()

  assert model_tape["users"].shape == (5000, 128)
  assert model_tape["items"].shape == (7500, 128)
  assert model_tape["ratings"] == ratings
  ```

  """
  tape_data = collections.OrderedDict({})

  def record(f, *args, **kwargs):
    """Records execution to a tape."""
    name = kwargs.get("name")
    output = interceptable(f)(*args, **kwargs)
    if name:
      tape_data[name] = output
    return output

  with interception(record):
    yield tape_data

def tape ( ) : tape_data = collections . OrderedDict ( { } ) def record ( f , * args , ** kwargs ) : name = kwargs . get ( "name" ) output = interceptable ( f ) ( * args , ** kwargs ) if name : tape_data [ name ] = output return output with interception ( record ) : yield tape_data

Context manager for recording interceptable executions onto a tape .


def toy_logistic_data(num_examples, input_size=2, weights_prior_stddev=5.0):
  """Generates synthetic data for binary classification.

  Args:
    num_examples: The number of samples to generate (scalar Python `int`).
    input_size: The input space dimension (scalar Python `int`).
    weights_prior_stddev: The prior standard deviation of the weight
      vector. (scalar Python `float`).

  Returns:
    random_weights: Sampled weights as a Numpy `array` of shape
      `[input_size]`.
    random_bias: Sampled bias as a scalar Python `float`.
    design_matrix: Points sampled uniformly from the cube `[-1,
       1]^{input_size}`, as a Numpy `array` of shape `(num_examples,
       input_size)`.
    labels: Labels sampled from the logistic model `p(label=1) =
      logistic(dot(features, random_weights) + random_bias)`, as a Numpy
      `int32` `array` of shape `(num_examples, 1)`.
  """
  random_weights = weights_prior_stddev * np.random.randn(input_size)
  random_bias = np.random.randn()
  design_matrix = np.random.rand(num_examples, input_size) * 2 - 1
  logits = np.reshape(
      np.dot(design_matrix, random_weights) + random_bias,
      (-1, 1))
  p_labels = 1. / (1 + np.exp(-logits))
  labels = np.int32(p_labels > np.random.rand(num_examples, 1))
  return random_weights, random_bias, np.float32(design_matrix), labels

def toy_logistic_data ( num_examples , input_size = 2 , weights_prior_stddev = 5.0 ) : random_weights = weights_prior_stddev * np . random . randn ( input_size ) random_bias = np . random . randn ( ) design_matrix = np . random . rand ( num_examples , input_size ) * 2 - 1 logits = np . reshape ( np . dot ( design_matrix , random_weights ) + random_bias , ( - 1 , 1 ) ) p_labels = 1. / ( 1 + np . exp ( - logits ) ) labels = np . int32 ( p_labels > np . random . rand ( num_examples , 1 ) ) return random_weights , random_bias , np . float32 ( design_matrix ) , labels

Generates synthetic data for binary classification .


def visualize_decision(features, labels, true_w_b, candidate_w_bs, fname):
  """Utility method to visualize decision boundaries in R^2.

  Args:
    features: Input points, as a Numpy `array` of shape `[num_examples, 2]`.
    labels: Numpy `float`-like array of shape `[num_examples, 1]` giving a
      label for each point.
    true_w_b: A `tuple` `(w, b)` where `w` is a Numpy array of
       shape `[2]` and `b` is a scalar `float`, interpreted as a
       decision rule of the form `dot(features, w) + b > 0`.
    candidate_w_bs: Python `iterable` containing tuples of the same form as
       true_w_b.
    fname: The filename to save the plot as a PNG image (Python `str`).
  """
  fig = figure.Figure(figsize=(6, 6))
  canvas = backend_agg.FigureCanvasAgg(fig)
  ax = fig.add_subplot(1, 1, 1)
  ax.scatter(features[:, 0], features[:, 1],
             c=np.float32(labels[:, 0]),
             cmap=cm.get_cmap("binary"),
             edgecolors="k")

  def plot_weights(w, b, **kwargs):
    w1, w2 = w
    x1s = np.linspace(-1, 1, 100)
    x2s = -(w1  * x1s + b) / w2
    ax.plot(x1s, x2s, **kwargs)

  for w, b in candidate_w_bs:
    plot_weights(w, b,
                 alpha=1./np.sqrt(len(candidate_w_bs)),
                 lw=1, color="blue")

  if true_w_b is not None:
    plot_weights(*true_w_b, lw=4,
                 color="green", label="true separator")

  ax.set_xlim([-1.5, 1.5])
  ax.set_ylim([-1.5, 1.5])
  ax.legend()

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

def
visualize_decision
(
features
,
labels
,
true_w_b
,
candidate_w_bs
,
fname
)
:
fig
=
figure
.
Figure
(
figsize
=
(
6
,
6
)
)
canvas
=
backend_agg
.
FigureCanvasAgg
(
fig
)
ax
=
fig
.
add_subplot
(
1
,
1
,
1
)
ax
.
scatter
(
features
[
:
,
0
]
,
features
[
:
,
1
]
,
c
=
np
.
float32
(
labels
[
:
,
0
]
)
,
cmap
=
cm
.
get_cmap
(
"binary"
)
,
edgecolors
=
"k"
)
def
plot_weights
(
w
,
b
,
**
kwargs
)
:
w1
,
w2
=
w
x1s
=
np
.
linspace
(
1
,
1
,
100
)
x2s
=
(
w1
*
x1s
+
b
)
/
w2
ax
.
plot
(
x1s
,
x2s
,
**
kwargs
)
for
w
,
b
in
candidate_w_bs
:
plot_weights
(
w
,
b
,
alpha
=
1.
/
np
.
sqrt
(
len
(
candidate_w_bs
)
)
,
lw
=
1
,
color
=
"blue"
)
if
true_w_b
is
not
:
plot_weights
(
*
true_w_b
,
lw
=
4
,
color
=
"green"
,
label
=
"true separator"
)
ax
.
set_xlim
(
[
1.5
,
1.5
]
)
ax
.
set_ylim
(
[
1.5
,
1.5
]
)
ax
.
legend
(
)
canvas
.
print_figure
(
fname
,
format
=
"png"
)
print
(
"saved {}"
.
format
(
fname
)
)

Utility method to visualize decision boundaries in R^2 .


def build_input_pipeline(x, y, batch_size):
  """Build a Dataset iterator for supervised classification.

  Args:
    x: Numpy `array` of features, indexed by the first dimension.
    y: Numpy `array` of labels, with the same first dimension as `x`.
    batch_size: Number of elements in each training batch.

  Returns:
    batch_features: `Tensor` feed  features, of shape
      `[batch_size] + x.shape[1:]`.
    batch_labels: `Tensor` feed of labels, of shape
      `[batch_size] + y.shape[1:]`.
  """
  training_dataset = tf.data.Dataset.from_tensor_slices((x, y))
  training_batches = training_dataset.repeat().batch(batch_size)
  training_iterator = tf.compat.v1.data.make_one_shot_iterator(training_batches)
  batch_features, batch_labels = training_iterator.get_next()
  return batch_features, batch_labels

def build_input_pipeline ( x , y , batch_size ) : training_dataset = tf . data . Dataset . from_tensor_slices ( ( x , y ) ) training_batches = training_dataset . repeat ( ) . batch ( batch_size ) training_iterator = tf . compat . v1 . data . make_one_shot_iterator ( training_batches ) batch_features , batch_labels = training_iterator . get_next ( ) return batch_features , batch_labels

Build a Dataset iterator for supervised classification .


def _maybe_check_valid_map_values(map_values, validate_args):
  """Validate `map_values` if `validate_args`==True."""
  assertions = []

  message = 'Rank of map_values must be 1.'
  if tensorshape_util.rank(map_values.shape) is not None:
    if tensorshape_util.rank(map_values.shape) != 1:
      raise ValueError(message)
  elif validate_args:
    assertions.append(assert_util.assert_rank(map_values, 1, message=message))

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

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

  return assertions

Validate map_values if validate_args == True .


def trace(state: State, fn: TransitionOperator, num_steps: IntTensor,
          trace_fn: Callable[[State, TensorNest], TensorNest]
         ) -> Tuple[State, TensorNest]:
  """`TransitionOperator` that runs `fn` repeatedly and traces its outputs.

  Args:
    state: A nest of `Tensor`s or None.
    fn: A `TransitionOperator`.
    num_steps: Number of steps to run the function for. Must be greater than 1.
    trace_fn: Callable that the unpacked outputs of `fn` and returns a nest of
      `Tensor`s. These will be stacked and returned.

  Returns:
    state: The final state returned by `fn`.
    traces: Stacked outputs of `trace_fn`.
  """

  def fn_wrapper(args, _):
    return tf.nest.map_structure(tf.convert_to_tensor, call_fn(fn, args[0]))

  def trace_fn_wrapper(args):
    return tf.nest.map_structure(tf.convert_to_tensor, call_fn(trace_fn, args))

  state = call_fn(fn, state)
  first_trace = trace_fn_wrapper(state)

  state, full_trace = mcmc_util.trace_scan(
      fn_wrapper, state, tf.ones(num_steps - 1), trace_fn=trace_fn_wrapper)

  prepend = lambda x, y: tf.concat(  # pylint: disable=g-long-lambda
      [tf.convert_to_tensor(value=x)[tf.newaxis], y], 0)

  return state, tf.nest.map_structure(prepend, first_trace, full_trace)

def trace ( state : State , fn : TransitionOperator , num_steps : IntTensor , trace_fn : Callable [ [ State , TensorNest ] , TensorNest ] ) -> Tuple [ State , TensorNest ] : def fn_wrapper ( args , _ ) : return tf . nest . map_structure ( tf . convert_to_tensor , call_fn ( fn , args [ 0 ] ) ) def trace_fn_wrapper ( args ) : return tf . nest . map_structure ( tf . convert_to_tensor , call_fn ( trace_fn , args ) ) state = call_fn ( fn , state ) first_trace = trace_fn_wrapper ( state ) state , full_trace = mcmc_util . trace_scan ( fn_wrapper , state , tf . ones ( num_steps - 1 ) , trace_fn = trace_fn_wrapper ) prepend = lambda x , y : tf . concat ( [ tf . convert_to_tensor ( value = x ) [ tf . newaxis ] , y ] , 0 ) return state , tf . nest . map_structure ( prepend , first_trace , full_trace )

TransitionOperator that runs fn repeatedly and traces its outputs .


def call_fn(fn: TransitionOperator, args: Union[Tuple[Any], Any]) -> Any:
  """Calls a transition operator with args, unpacking args if its a sequence.

  Args:
    fn: A `TransitionOperator`.
    args: Arguments to `fn`

  Returns:
    ret: Return value of `fn`.
  """

  if isinstance(args, (list, tuple)) and not mcmc_util.is_namedtuple_like(args):
    args = args  # type: Tuple[Any]
    return fn(*args)
  else:
    return fn(args)

def call_fn ( fn : TransitionOperator , args : Union [ Tuple [ Any ] , Any ] ) -> Any : if isinstance ( args , ( list , tuple ) ) and not mcmc_util . is_namedtuple_like ( args ) : args = args return fn ( * args ) else : return fn ( args )

Calls a transition operator with args unpacking args if its a sequence .


def call_and_grads(fn: TransitionOperator, args: Union[Tuple[Any], Any]
                  ) -> Tuple[tf.Tensor, TensorNest, TensorNest]:
  """Calls `fn` and returns the gradients with respect to `fn`'s first output.

  Args:
    fn: A `TransitionOperator`.
    args: Arguments to `fn`

  Returns:
    ret: First output of `fn`.
    extra: Second output of `fn`.
    grads: Gradients of `ret` with respect to `args`.
  """
  with tf.GradientTape() as tape:
    tape.watch(args)
    ret, extra = call_fn(fn, args)
  grads = tape.gradient(ret, args)
  return ret, extra, grads

def call_and_grads ( fn : TransitionOperator , args : Union [ Tuple [ Any ] , Any ] ) -> Tuple [ tf . Tensor , TensorNest , TensorNest ] : with tf . GradientTape ( ) as tape : tape . watch ( args ) ret , extra = call_fn ( fn , args ) grads = tape . gradient ( ret , args ) return ret , extra , grads

Calls fn and returns the gradients with respect to fn s first output .


def maybe_broadcast_structure(from_structure: Any, to_structure: Any) -> Any:
  """Maybe broadcasts `from_structure` to `to_structure`.

  If `from_structure` is a singleton, it is tiled to match the structure of
  `to_structure`. Note that the elements in `from_structure` are not copied if
  this tiling occurs.

  Args:
    from_structure: A structure.
    to_structure: A structure.

  Returns:
    new_from_structure: Same structure as `to_structure`.
  """
  flat_from = tf.nest.flatten(from_structure)
  flat_to = tf.nest.flatten(to_structure)
  if len(flat_from) == 1:
    flat_from *= len(flat_to)
  return tf.nest.pack_sequence_as(to_structure, flat_from)

def maybe_broadcast_structure ( from_structure : Any , to_structure : Any ) -> Any : flat_from = tf . nest . flatten ( from_structure ) flat_to = tf . nest . flatten ( to_structure ) if len ( flat_from ) == 1 : flat_from *= len ( flat_to ) return tf . nest . pack_sequence_as ( to_structure , flat_from )

Maybe broadcasts from_structure to to_structure .


def transform_log_prob_fn(log_prob_fn: PotentialFn,
                          bijector: BijectorNest,
                          init_state: State = None
                         ) -> Union[PotentialFn, Tuple[PotentialFn, State]]:
  """Transforms a log-prob function using a bijector.

  This takes a log-prob function and creates a new log-prob function that now
  takes takes state in the domain of the bijector, forward transforms that state
  and calls the original log-prob function. It then returns the log-probability
  that correctly accounts for this transformation.

  The forward-transformed state is pre-pended to the original log-prob
  function's extra returns and returned as the new extra return.

  For convenience you can also pass the initial state (in the original space),
  and this function will return the inverse transformed as the 2nd return value.
  You'd use this to initialize MCMC operators that operate in the transformed
  space.

  Args:
    log_prob_fn: Log prob fn.
    bijector: Bijector(s), must be of the same structure as the `log_prob_fn`
      inputs.
    init_state: Initial state, in the original space.

  Returns:
    transformed_log_prob_fn: Transformed log prob fn.
    transformed_init_state: If `init_state` is provided. Initial state in the
      transformed space.
  """

  def wrapper(*args):
    """Transformed wrapper."""
    bijector_ = bijector

    args = tf.nest.map_structure(lambda x: 0. + x, args)
    if len(args) == 1:
      args = args[0]
    elif isinstance(bijector_, list):
      bijector_ = tuple(bijector_)

    original_space_args = tf.nest.map_structure(lambda b, x: b.forward(x),
                                                bijector_, args)
    original_space_args = original_space_args  # type: Tuple[Any]
    original_space_log_prob, extra = call_fn(log_prob_fn, original_space_args)
    event_ndims = tf.nest.map_structure(
        lambda x: tf.rank(x) - tf.rank(original_space_log_prob), args)

    return original_space_log_prob + sum(
        tf.nest.flatten(
            tf.nest.map_structure(
                lambda b, x, e: b.forward_log_det_jacobian(x, event_ndims=e),
                bijector_, args, event_ndims))), [original_space_args, extra]

  if init_state is None:
    return wrapper
  else:
    return wrapper, tf.nest.map_structure(lambda b, s: b.inverse(s), bijector,
                                          init_state)

def transform_log_prob_fn ( log_prob_fn : PotentialFn , bijector : BijectorNest , init_state : State = None ) -> Union [ PotentialFn , Tuple [ PotentialFn , State ] ] : def wrapper ( * args ) : bijector_ = bijector args = tf . nest . map_structure ( lambda x : 0. + x , args ) if len ( args ) == 1 : args = args [ 0 ] elif isinstance ( bijector_ , list ) : bijector_ = tuple ( bijector_ ) original_space_args = tf . nest . map_structure ( lambda b , x : b . forward ( x ) , bijector_ , args ) original_space_args = original_space_args original_space_log_prob , extra = call_fn ( log_prob_fn , original_space_args ) event_ndims = tf . nest . map_structure ( lambda x : tf . rank ( x ) - tf . rank ( original_space_log_prob ) , args ) return original_space_log_prob + sum ( tf . nest . flatten ( tf . nest . map_structure ( lambda b , x , e : b . forward_log_det_jacobian ( x , event_ndims = e ) , bijector_ , args , event_ndims ) ) ) , [ original_space_args , extra ] if init_state is None : return wrapper else : return wrapper , tf . nest . map_structure ( lambda b , s : b . inverse ( s ) , bijector , init_state )

Transforms a log - prob function using a bijector .


def leapfrog_step(leapfrog_step_state: LeapFrogStepState,
                  step_size: FloatTensor, target_log_prob_fn: PotentialFn,
                  kinetic_energy_fn: PotentialFn
                 ) -> Tuple[LeapFrogStepState, LeapFrogStepExtras]:
  """Leapfrog `TransitionOperator`.

  Args:
    leapfrog_step_state: LeapFrogStepState.
    step_size: Step size, structure broadcastable to the `target_log_prob_fn`
      state.
    target_log_prob_fn: Target log prob fn.
    kinetic_energy_fn: Kinetic energy fn.

  Returns:
    leapfrog_step_state: LeapFrogStepState.
    leapfrog_step_extras: LeapFrogStepExtras.
  """
  state = leapfrog_step_state.state
  state_grads = leapfrog_step_state.state_grads
  momentum = leapfrog_step_state.momentum
  step_size = maybe_broadcast_structure(step_size, state)

  state = tf.nest.map_structure(tf.convert_to_tensor, state)
  momentum = tf.nest.map_structure(tf.convert_to_tensor, momentum)
  state = tf.nest.map_structure(tf.convert_to_tensor, state)

  if state_grads is None:
    _, _, state_grads = call_and_grads(target_log_prob_fn, state)
  else:
    state_grads = tf.nest.map_structure(tf.convert_to_tensor, state_grads)

  momentum = tf.nest.map_structure(lambda m, sg, s: m + 0.5 * sg * s, momentum,
                                   state_grads, step_size)

  kinetic_energy, kinetic_energy_extra, momentum_grads = call_and_grads(
      kinetic_energy_fn, momentum)

  state = tf.nest.map_structure(lambda x, mg, s: x + mg * s, state,
                                momentum_grads, step_size)

  target_log_prob, state_extra, state_grads = call_and_grads(
      target_log_prob_fn, state)

  momentum = tf.nest.map_structure(lambda m, sg, s: m + 0.5 * sg * s, momentum,
                                   state_grads, step_size)

  return LeapFrogStepState(state, state_grads, momentum), LeapFrogStepExtras(
      target_log_prob, state_extra, kinetic_energy, kinetic_energy_extra)

def leapfrog_step ( leapfrog_step_state : LeapFrogStepState , step_size : FloatTensor , target_log_prob_fn : PotentialFn , kinetic_energy_fn : PotentialFn ) -> Tuple [ LeapFrogStepState , LeapFrogStepExtras ] : state = leapfrog_step_state . state state_grads = leapfrog_step_state . state_grads momentum = leapfrog_step_state . momentum step_size = maybe_broadcast_structure ( step_size , state ) state = tf . nest . map_structure ( tf . convert_to_tensor , state ) momentum = tf . nest . map_structure ( tf . convert_to_tensor , momentum ) state = tf . nest . map_structure ( tf . convert_to_tensor , state ) if state_grads is None : _ , _ , state_grads = call_and_grads ( target_log_prob_fn , state ) else : state_grads = tf . nest . map_structure ( tf . convert_to_tensor , state_grads ) momentum = tf . nest . map_structure ( lambda m , sg , s : m + 0.5 * sg * s , momentum , state_grads , step_size ) kinetic_energy , kinetic_energy_extra , momentum_grads = call_and_grads ( kinetic_energy_fn , momentum ) state = tf . nest . map_structure ( lambda x , mg , s : x + mg * s , state , momentum_grads , step_size ) target_log_prob , state_extra , state_grads = call_and_grads ( target_log_prob_fn , state ) momentum = tf . nest . map_structure ( lambda m , sg , s : m + 0.5 * sg * s , momentum , state_grads , step_size ) return LeapFrogStepState ( state , state_grads , momentum ) , LeapFrogStepExtras ( target_log_prob , state_extra , kinetic_energy , kinetic_energy_extra )


def metropolis_hastings_step(current_state: State,
                             proposed_state: State,
                             energy_change: FloatTensor,
                             seed=None) -> Tuple[State, tf.Tensor, tf.Tensor]:
  """Metropolis-Hastings step.

  This probabilistically chooses between `current_state` and `proposed_state`
  based on the `energy_change` so as to preserve detailed balance.

  Energy change is the negative of `log_accept_ratio`.

  Args:
    current_state: Current state.
    proposed_state: Proposed state.
    energy_change: E(proposed_state) - E(previous_state).
    seed: For reproducibility.

  Returns:
    new_state: The chosen state.
    is_accepted: Whether the proposed state was accepted.
    log_uniform: The random number that was used to select between the two
      states.
  """
  flat_current = tf.nest.flatten(current_state)
  flat_proposed = nest.flatten_up_to(current_state, proposed_state)
  # Impute the None's in the current state.
  flat_current = [
      p if c is None else c for p, c in zip(flat_proposed, flat_current)
  ]
  current_state = tf.nest.pack_sequence_as(current_state, flat_current)

  current_state = tf.nest.map_structure(tf.convert_to_tensor, current_state)
  proposed_state = tf.nest.map_structure(tf.convert_to_tensor, proposed_state)
  energy_change = tf.convert_to_tensor(value=energy_change)

  log_accept_ratio = -energy_change

  log_uniform = tf.math.log(
      tf.random.uniform(
          shape=tf.shape(input=log_accept_ratio),
          dtype=log_accept_ratio.dtype.base_dtype,
          seed=seed))
  is_accepted = log_uniform < log_accept_ratio

  next_state = mcmc_util.choose(
      is_accepted, proposed_state, current_state, name='choose_next_state')
  return next_state, is_accepted, log_uniform

def metropolis_hastings_step ( current_state : State , proposed_state : State , energy_change : FloatTensor , seed = None ) -> Tuple [ State , tf . Tensor , tf . Tensor ] : flat_current = tf . nest . flatten ( current_state ) flat_proposed = nest . flatten_up_to ( current_state , proposed_state ) flat_current = [ p if c is None else c for p , c in zip ( flat_proposed , flat_current ) ] current_state = tf . nest . pack_sequence_as ( current_state , flat_current ) current_state = tf . nest . map_structure ( tf . convert_to_tensor , current_state ) proposed_state = tf . nest . map_structure ( tf . convert_to_tensor , proposed_state ) energy_change = tf . convert_to_tensor ( value = energy_change ) log_accept_ratio = - energy_change log_uniform = tf . math . log ( tf . random . uniform ( shape = tf . shape ( input = log_accept_ratio ) , dtype = log_accept_ratio . dtype . base_dtype , seed = seed ) ) is_accepted = log_uniform < log_accept_ratio next_state = mcmc_util . choose ( is_accepted , proposed_state , current_state , name = 'choose_next_state' ) return next_state , is_accepted , log_uniform


def hamiltonian_monte_carlo(
    hmc_state: HamiltonianMonteCarloState,
    target_log_prob_fn: PotentialFn,
    step_size: Any,
    num_leapfrog_steps: IntTensor,
    momentum: State = None,
    kinetic_energy_fn: PotentialFn = None,
    momentum_sample_fn: MomentumSampleFn = None,
    leapfrog_trace_fn: Callable[[LeapFrogStepState, LeapFrogStepExtras],
                                TensorNest] = lambda *args: (),
    seed=None,
) -> Tuple[HamiltonianMonteCarloState, HamiltonianMonteCarloExtra]:
  """Hamiltonian Monte Carlo `TransitionOperator`.

  #### Example

  ```python
  step_size = 0.2
  num_steps = 2000
  num_leapfrog_steps = 10
  state = tf.ones([16, 2])

  base_mean = [1., 0]
  base_cov = [[1, 0.5], [0.5, 1]]

  bijector = tfb.Softplus()
  base_dist = tfd.MultivariateNormalFullCovariance(
      loc=base_mean, covariance_matrix=base_cov)
  target_dist = bijector(base_dist)

  def orig_target_log_prob_fn(x):
    return target_dist.log_prob(x), ()

  target_log_prob_fn, state = fun_mcmc.transform_log_prob_fn(
      orig_target_log_prob_fn, bijector, state)

  kernel = tf.function(lambda state: fun_mcmc.hamiltonian_monte_carlo(
      state,
      step_size=step_size,
      num_leapfrog_steps=num_leapfrog_steps,
      target_log_prob_fn=target_log_prob_fn,
      seed=tfp_test_util.test_seed()))

  _, chain = fun_mcmc.trace(
      state=fun_mcmc.HamiltonianMonteCarloState(
          state=state,
          state_grads=None,
          target_log_prob=None,
          state_extra=None),
      fn=kernel,
      num_steps=num_steps,
      trace_fn=lambda state, extra: state.state_extra[0])
  ```

  Args:
    hmc_state: HamiltonianMonteCarloState.
    target_log_prob_fn: Target log prob fn.
    step_size: Step size, structure broadcastable to the `target_log_prob_fn`
      state.
    num_leapfrog_steps: Number of leapfrog steps to take.
    momentum: Initial momentum, passed to `momentum_sample_fn`. Default: zeroes.
    kinetic_energy_fn: Kinetic energy function.
    momentum_sample_fn: Sampler for the momentum.
    leapfrog_trace_fn: Trace function for the leapfrog integrator.
    seed: For reproducibility.

  Returns:
    hmc_state: HamiltonianMonteCarloState
    hmc_extra: HamiltonianMonteCarloExtra
  """
  state = hmc_state.state
  state_grads = hmc_state.state_grads
  target_log_prob = hmc_state.target_log_prob
  state_extra = hmc_state.state_extra

  if kinetic_energy_fn is None:

    # pylint: disable=function-redefined
    def kinetic_energy_fn(*momentum):
      return tf.add_n([
          tf.reduce_sum(input_tensor=tf.square(x), axis=-1) / 2.
          for x in tf.nest.flatten(momentum)
      ]), ()

  if momentum_sample_fn is None:

    # pylint: disable=function-redefined
    def momentum_sample_fn(*momentum):
      ret = tf.nest.map_structure(
          lambda x: tf.random.normal(tf.shape(input=x), dtype=x.dtype),
          momentum)
      if len(ret) == 1:
        return ret[0]
      else:
        return ret

  if momentum is None:
    momentum = call_fn(momentum_sample_fn,
                       tf.nest.map_structure(tf.zeros_like, state))
  if target_log_prob is None:
    target_log_prob, state_extra, state_grads = call_and_grads(
        target_log_prob_fn, state)

  kinetic_energy, _ = call_fn(kinetic_energy_fn, momentum)
  current_energy = -target_log_prob + kinetic_energy
  current_state = HamiltonianMonteCarloState(
      state=state,
      state_grads=state_grads,
      state_extra=state_extra,
      target_log_prob=target_log_prob)

  def leapfrog_wrapper(leapfrog_state, target_log_prob, state_extra):
    """Leapfrog wrapper that tracks extra state."""
    del target_log_prob
    del state_extra

    leapfrog_state, leapfrog_extra = leapfrog_step(
        leapfrog_state,
        step_size=step_size,
        target_log_prob_fn=target_log_prob_fn,
        kinetic_energy_fn=kinetic_energy_fn)

    return [
        leapfrog_state, leapfrog_extra.target_log_prob,
        leapfrog_extra.state_extra
    ], leapfrog_extra

  def leapfrog_trace_wrapper_fn(args, leapfrog_extra):
    return leapfrog_trace_fn(args[0], leapfrog_extra)

  leapfrog_wrapper_state = (LeapFrogStepState(state, state_grads, momentum),
                            target_log_prob, state_extra)

  [[leapfrog_state, target_log_prob, state_extra], _], leapfrog_trace = trace(
      leapfrog_wrapper_state,
      leapfrog_wrapper,
      num_leapfrog_steps,
      trace_fn=leapfrog_trace_wrapper_fn)

  kinetic_energy, _ = call_fn(kinetic_energy_fn, leapfrog_state.momentum)
  proposed_energy = -target_log_prob + kinetic_energy
  proposed_state = HamiltonianMonteCarloState(
      state=leapfrog_state.state,
      state_grads=leapfrog_state.state_grads,
      target_log_prob=target_log_prob,
      state_extra=state_extra)

  energy_change = proposed_energy - current_energy
  hmc_state, is_accepted, _ = metropolis_hastings_step(
      current_state, proposed_state, energy_change, seed=seed)

  hmc_state = hmc_state  # type: HamiltonianMonteCarloState
  return hmc_state, HamiltonianMonteCarloExtra(
      is_accepted=is_accepted,
      proposed_hmc_state=proposed_state,
      log_accept_ratio=-energy_change,
      leapfrog_trace=leapfrog_trace)

def hamiltonian_monte_carlo ( hmc_state : HamiltonianMonteCarloState , target_log_prob_fn : PotentialFn , step_size : Any , num_leapfrog_steps : IntTensor , momentum : State = None , kinetic_energy_fn : PotentialFn = None , momentum_sample_fn : MomentumSampleFn = None , leapfrog_trace_fn : Callable [ [ LeapFrogStepState , LeapFrogStepExtras ] , TensorNest ] = lambda * args : ( ) , seed = None , ) -> Tuple [ HamiltonianMonteCarloState , HamiltonianMonteCarloExtra ] : state = hmc_state . state state_grads = hmc_state . state_grads target_log_prob = hmc_state . target_log_prob state_extra = hmc_state . state_extra if kinetic_energy_fn is None : def kinetic_energy_fn ( * momentum ) : return tf . add_n ( [ tf . reduce_sum ( input_tensor = tf . square ( x ) , axis = - 1 ) / 2. for x in tf . nest . flatten ( momentum ) ] ) , ( ) if momentum_sample_fn is None : def momentum_sample_fn ( * momentum ) : ret = tf . nest . map_structure ( lambda x : tf . random . normal ( tf . shape ( input = x ) , dtype = x . dtype ) , momentum ) if len ( ret ) == 1 : return ret [ 0 ] else : return ret if momentum is None : momentum = call_fn ( momentum_sample_fn , tf . nest . map_structure ( tf . zeros_like , state ) ) if target_log_prob is None : target_log_prob , state_extra , state_grads = call_and_grads ( target_log_prob_fn , state ) kinetic_energy , _ = call_fn ( kinetic_energy_fn , momentum ) current_energy = - target_log_prob + kinetic_energy current_state = HamiltonianMonteCarloState ( state = state , state_grads = state_grads , state_extra = state_extra , target_log_prob = target_log_prob ) def leapfrog_wrapper ( leapfrog_state , target_log_prob , state_extra ) : del target_log_prob del state_extra leapfrog_state , leapfrog_extra = leapfrog_step ( leapfrog_state , step_size = step_size , target_log_prob_fn = target_log_prob_fn , kinetic_energy_fn = kinetic_energy_fn ) return [ leapfrog_state , leapfrog_extra . target_log_prob , leapfrog_extra . state_extra ] , leapfrog_extra def leapfrog_trace_wrapper_fn ( args , leapfrog_extra ) : return leapfrog_trace_fn ( args [ 0 ] , leapfrog_extra ) leapfrog_wrapper_state = ( LeapFrogStepState ( state , state_grads , momentum ) , target_log_prob , state_extra ) [ [ leapfrog_state , target_log_prob , state_extra ] , _ ] , leapfrog_trace = trace ( leapfrog_wrapper_state , leapfrog_wrapper , num_leapfrog_steps , trace_fn = leapfrog_trace_wrapper_fn ) kinetic_energy , _ = call_fn ( kinetic_energy_fn , leapfrog_state . momentum ) proposed_energy = - target_log_prob + kinetic_energy proposed_state = HamiltonianMonteCarloState ( state = leapfrog_state . state , state_grads = leapfrog_state . state_grads , target_log_prob = target_log_prob , state_extra = state_extra ) energy_change = proposed_energy - current_energy hmc_state , is_accepted , _ = metropolis_hastings_step ( current_state , proposed_state , energy_change , seed = seed ) hmc_state = hmc_state return hmc_state , HamiltonianMonteCarloExtra ( is_accepted = is_accepted , proposed_hmc_state = proposed_state , log_accept_ratio = - energy_change , leapfrog_trace = leapfrog_trace )

Hamiltonian Monte Carlo TransitionOperator .


def sign_adaptation(control: FloatNest,
                    output: FloatTensor,
                    set_point: FloatTensor,
                    adaptation_rate: FloatTensor = 0.01) -> FloatNest:
  """A function to do simple sign-based control of a variable.

  ```
  control = control * (1. + adaptation_rate) ** sign(output - set_point)
  ```

  Args:
    control: The control variable.
    output: The output variable.
    set_point: The set point for `output`. This function will adjust `control`
      so that `output` matches `set_point`.
    adaptation_rate: Adaptation rate.

  Returns:
    control: New control.
  """

  def _get_new_control(control, output, set_point):
    new_control = mcmc_util.choose(output > set_point,
                                   control * (1. + adaptation_rate),
                                   control / (1. + adaptation_rate))
    return new_control

  output = maybe_broadcast_structure(output, control)
  set_point = maybe_broadcast_structure(set_point, control)

  return tf.nest.map_structure(_get_new_control, control, output, set_point)

def sign_adaptation ( control : FloatNest , output : FloatTensor , set_point : FloatTensor , adaptation_rate : FloatTensor = 0.01 ) -> FloatNest : def _get_new_control ( control , output , set_point ) : new_control = mcmc_util . choose ( output > set_point , control * ( 1. + adaptation_rate ) , control / ( 1. + adaptation_rate ) ) return new_control output = maybe_broadcast_structure ( output , control ) set_point = maybe_broadcast_structure ( set_point , control ) return tf . nest . map_structure ( _get_new_control , control , output , set_point )

A function to do simple sign - based control of a variable .


def from_config(cls, config):
    """Creates a layer from its config.

    This method is the reverse of `get_config`, capable of instantiating the
    same layer from the config dictionary.

    Args:
      config: A Python dictionary, typically the output of `get_config`.

    Returns:
      layer: A layer instance.
    """
    config = config.copy()
    function_keys = [
        'kernel_posterior_fn',
        'kernel_posterior_tensor_fn',
        'kernel_prior_fn',
        'kernel_divergence_fn',
        'bias_posterior_fn',
        'bias_posterior_tensor_fn',
        'bias_prior_fn',
        'bias_divergence_fn',
    ]
    for function_key in function_keys:
      serial = config[function_key]
      function_type = config.pop(function_key + '_type')
      if serial is not None:
        config[function_key] = tfp_layers_util.deserialize_function(
            serial,
            function_type=function_type)
    return cls(**config)

def from_config ( cls , config ) : config = config . copy ( ) function_keys = [ 'kernel_posterior_fn' , 'kernel_posterior_tensor_fn' , 'kernel_prior_fn' , 'kernel_divergence_fn' , 'bias_posterior_fn' , 'bias_posterior_tensor_fn' , 'bias_prior_fn' , 'bias_divergence_fn' , ] for function_key in function_keys : serial = config [ function_key ] function_type = config . pop ( function_key + '_type' ) if serial is not None : config [ function_key ] = tfp_layers_util . deserialize_function ( serial , function_type = function_type ) return cls ( ** config )


def _as_tensor(x, name, dtype):
  """Convenience to convert to `Tensor` or leave as `None`."""
  return None if x is None else tf.convert_to_tensor(
      value=x, name=name, dtype=dtype)

def _as_tensor ( x , name , dtype ) : return None if x is None else tf . convert_to_tensor ( value = x , name = name , dtype = dtype )

Convenience to convert to Tensor or leave as None .


def _create_scale_operator(self, identity_multiplier, diag, tril,
                             perturb_diag, perturb_factor, shift, validate_args,
                             dtype):
    """Construct `scale` from various components.

    Args:
      identity_multiplier: floating point rank 0 `Tensor` representing a scaling
        done to the identity matrix.
      diag: Floating-point `Tensor` representing the diagonal matrix.`diag` has
        shape `[N1, N2, ...  k]`, which represents a k x k diagonal matrix.
      tril: Floating-point `Tensor` representing the lower triangular matrix.
       `tril` has shape `[N1, N2, ...  k, k]`, which represents a k x k lower
       triangular matrix.
      perturb_diag: Floating-point `Tensor` representing the diagonal matrix of
        the low rank update.
      perturb_factor: Floating-point `Tensor` representing factor matrix.
      shift: Floating-point `Tensor` representing `shift in `scale @ X + shift`.
      validate_args: Python `bool` indicating whether arguments should be
        checked for correctness.
      dtype: `DType` for arg `Tensor` conversions.

    Returns:
      scale. In the case of scaling by a constant, scale is a
      floating point `Tensor`. Otherwise, scale is a `LinearOperator`.

    Raises:
      ValueError: if all of `tril`, `diag` and `identity_multiplier` are `None`.
    """
    identity_multiplier = _as_tensor(identity_multiplier, "identity_multiplier",
                                     dtype)
    diag = _as_tensor(diag, "diag", dtype)
    tril = _as_tensor(tril, "tril", dtype)
    perturb_diag = _as_tensor(perturb_diag, "perturb_diag", dtype)
    perturb_factor = _as_tensor(perturb_factor, "perturb_factor", dtype)

    # If possible, use the low rank update to infer the shape of
    # the identity matrix, when scale represents a scaled identity matrix
    # with a low rank update.
    shape_hint = None
    if perturb_factor is not None:
      shape_hint = distribution_util.dimension_size(perturb_factor, axis=-2)

    if self._is_only_identity_multiplier:
      if validate_args:
        return distribution_util.with_dependencies([
            assert_util.assert_none_equal(
                identity_multiplier, tf.zeros([], identity_multiplier.dtype),
                ["identity_multiplier should be non-zero."])
        ], identity_multiplier)
      return identity_multiplier

    scale = distribution_util.make_tril_scale(
        loc=shift,
        scale_tril=tril,
        scale_diag=diag,
        scale_identity_multiplier=identity_multiplier,
        validate_args=validate_args,
        assert_positive=False,
        shape_hint=shape_hint)

    if perturb_factor is not None:
      return tf.linalg.LinearOperatorLowRankUpdate(
          scale,
          u=perturb_factor,
          diag_update=perturb_diag,
          is_diag_update_positive=perturb_diag is None,
          is_non_singular=True,  # Implied by is_positive_definite=True.
          is_self_adjoint=True,
          is_positive_definite=True,
          is_square=True)

    return scale

Construct scale from various components .


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

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

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

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

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

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

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

      return next_state_parts
  return _fn

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


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

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

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

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

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

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

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

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


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

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

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

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

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


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

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

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

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

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

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

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

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

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

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

def entropy_lower_bound ( self , name = "entropy_lower_bound" ) : r with self . _name_scope ( name ) : with tf . control_dependencies ( self . _assertions ) : distribution_entropies = [ d . entropy ( ) for d in self . components ] cat_probs = self . _cat_probs ( log_probs = False ) partial_entropies = [ c_p * m for ( c_p , m ) in zip ( cat_probs , distribution_entropies ) ] return tf . add_n ( partial_entropies )

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


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

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

Get a list of num_components batchwise probabilities .


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

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

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

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

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

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

  return assertions

Validate outcomes logits and probs s shapes .


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

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

  import distutils.version

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

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

Attempt to import tensorflow and ensure its version is sufficient .


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

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

Bayesian logistic regression which returns labels given features .


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

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

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

def covertype ( ) : import sklearn . datasets data = sklearn . datasets . covtype . fetch_covtype ( ) features = data . data labels = data . target features -= features . mean ( 0 ) features /= features . std ( 0 ) features = np . hstack ( [ features , np . ones ( [ features . shape [ 0 ] , 1 ] ) ] ) features = tf . cast ( features , dtype = tf . float32 ) _ , counts = np . unique ( labels , return_counts = True ) specific_category = np . argmax ( counts ) labels = ( labels == specific_category ) labels = tf . cast ( labels , dtype = tf . int32 ) return features , labels


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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


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

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

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

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

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

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

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

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

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

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

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

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

Estimate standard deviation using samples .


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

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

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

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

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

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

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

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

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


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

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

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

  return positive_axis

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

Rectify possibly negatively axis . Prefer return Python list .


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

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

A version of squeeze that works with dynamic axis .


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

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


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

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

Reconstruct input x from a its normalized version .


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

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

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

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


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

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

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

def semilocal_linear_trend_transition_noise ( level_scale , slope_mean , slope_scale , autoregressive_coef ) : broadcast_batch_shape = dist_util . get_broadcast_shape ( level_scale , slope_mean , slope_scale , autoregressive_coef ) broadcast_ones = tf . ones ( broadcast_batch_shape , dtype = level_scale . dtype ) scale_diag = tf . stack ( [ level_scale * broadcast_ones , slope_scale * broadcast_ones ] , axis = - 1 ) bias = tf . stack ( [ tf . zeros_like ( broadcast_ones ) , slope_mean * ( 1 - autoregressive_coef ) * broadcast_ones ] , axis = - 1 ) return tfd . MultivariateNormalDiag ( loc = bias , scale_diag = scale_diag )

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


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

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

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

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

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

  #### Examples

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

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

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

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

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

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

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

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

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

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

  #### References

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

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

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

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

    max_size = tf.reduce_max(input_tensor=max_sizes_by_axes)

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

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

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

r Returns a sample from the dim dimensional Halton sequence .