def build_factored_variational_loss(model,
                                    observed_time_series,
                                    init_batch_shape=(),
                                    seed=None,
                                    name=None):
  """Build a loss function for variational inference in STS models.

  Variational inference searches for the distribution within some family of
  approximate posteriors that minimizes a divergence between the approximate
  posterior `q(z)` and true posterior `p(z|observed_time_series)`. By converting
  inference to optimization, it's generally much faster than sampling-based
  inference algorithms such as HMC. The tradeoff is that the approximating
  family rarely contains the true posterior, so it may miss important aspects of
  posterior structure (in particular, dependence between variables) and should
  not be blindly trusted. Results may vary; it's generally wise to compare to
  HMC to evaluate whether inference quality is sufficient for your task at hand.

  This method constructs a loss function for variational inference using the
  Kullback-Liebler divergence `KL[q(z) || p(z|observed_time_series)]`, with an
  approximating family given by independent Normal distributions transformed to
  the appropriate parameter space for each parameter. Minimizing this loss (the
  negative ELBO) maximizes a lower bound on the log model evidence `-log
  p(observed_time_series)`. This is equivalent to the 'mean-field' method
  implemented in [1]. and is a standard approach. The resulting posterior
  approximations are unimodal; they will tend to underestimate posterior
  uncertainty when the true posterior contains multiple modes (the `KL[q||p]`
  divergence encourages choosing a single mode) or dependence between variables.

  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`, which includes
      a mask `Tensor` to specify timesteps with missing observations.
    init_batch_shape: Batch shape (Python `tuple`, `list`, or `int`) of initial
      states to optimize in parallel.
      Default value: `()`. (i.e., just run a single optimization).
    seed: Python integer to seed the random number generator.
    name: Python `str` name prefixed to ops created by this function.
      Default value: `None` (i.e., 'build_factored_variational_loss').

  Returns:
    variational_loss: `float` `Tensor` of shape
      `concat([init_batch_shape, model.batch_shape])`, encoding a stochastic
      estimate of an upper bound on the negative model evidence `-log p(y)`.
      Minimizing this loss performs variational inference; the gap between the
      variational bound and the true (generally unknown) model evidence
      corresponds to the divergence `KL[q||p]` between the approximate and true
      posterior.
    variational_distributions: `collections.OrderedDict` giving
      the approximate posterior for each model parameter. The keys are
      Python `str` parameter names in order, corresponding to
      `[param.name for param in model.parameters]`. The values are
      `tfd.Distribution` instances with batch shape
      `concat([init_batch_shape, model.batch_shape])`; these will typically be
      of the form `tfd.TransformedDistribution(tfd.Normal(...),
      bijector=param.bijector)`.

  #### Examples

  Assume we've built a structural time-series model:

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

  To run variational inference, we simply construct the loss and optimize
  it:

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

    train_op = tf.train.AdamOptimizer(0.1).minimize(variational_loss)
    with tf.Session() as sess:
      sess.run(tf.global_variables_initializer())

      for step in range(200):
        _, loss_ = sess.run((train_op, variational_loss))

        if step % 20 == 0:
          print("step {} loss {}".format(step, loss_))

      posterior_samples_ = sess.run({
        param_name: q.sample(50)
        for param_name, q in variational_distributions.items()})
  ```

  As a more complex example, we might try to avoid local optima by optimizing
  from multiple initializations in parallel, and selecting the result with the
  lowest loss:

  ```python
    (variational_loss,
     variational_distributions) = tfp.sts.build_factored_variational_loss(
       model=model, observed_time_series=observed_time_series,
       init_batch_shape=[10])

    train_op = tf.train.AdamOptimizer(0.1).minimize(variational_loss)
    with tf.Session() as sess:
      sess.run(tf.global_variables_initializer())

      for step in range(200):
        _, loss_ = sess.run((train_op, variational_loss))

        if step % 20 == 0:
          print("step {} losses {}".format(step, loss_))

      # Draw multiple samples to reduce Monte Carlo error in the optimized
      # variational bounds.
      avg_loss = np.mean(
        [sess.run(variational_loss) for _ in range(25)], axis=0)
      best_posterior_idx = np.argmin(avg_loss, axis=0).astype(np.int32)
  ```

  #### References

  [1]: Alp Kucukelbir, Dustin Tran, Rajesh Ranganath, Andrew Gelman, and
       David M. Blei. Automatic Differentiation Variational Inference. In
       _Journal of Machine Learning Research_, 2017.
       https://arxiv.org/abs/1603.00788

  """

  with tf.compat.v1.name_scope(
      name, 'build_factored_variational_loss',
      values=[observed_time_series]) as name:
    seed = tfd.SeedStream(
        seed, salt='StructuralTimeSeries_build_factored_variational_loss')

    variational_distributions = collections.OrderedDict()
    variational_samples = []
    for param in model.parameters:
      def initial_loc_fn(param):
        return sample_uniform_initial_state(
            param, return_constrained=True,
            init_sample_shape=init_batch_shape,
            seed=seed())
      q = _build_trainable_posterior(param, initial_loc_fn=initial_loc_fn)
      variational_distributions[param.name] = q
      variational_samples.append(q.sample(seed=seed()))

    # Multiple initializations (similar to HMC chains) manifest as an extra
    # param batch dimension, so we need to add corresponding batch dimension(s)
    # to `observed_time_series`.
    observed_time_series = sts_util.pad_batch_dimension_for_multiple_chains(
        observed_time_series, model, chain_batch_shape=init_batch_shape)

    # Construct the variational bound.
    log_prob_fn = model.joint_log_prob(observed_time_series)
    expected_log_joint = log_prob_fn(*variational_samples)
    entropy = tf.reduce_sum(
        input_tensor=[
            -q.log_prob(sample) for (q, sample) in zip(
                variational_distributions.values(), variational_samples)
        ],
        axis=0)
    variational_loss = -(expected_log_joint + entropy)  # -ELBO

  return variational_loss, variational_distributions

Build a loss function for variational inference in STS models .


def _minimize_in_graph(build_loss_fn, num_steps=200, optimizer=None):
  """Run an optimizer within the graph to minimize a loss function."""
  optimizer = tf.compat.v1.train.AdamOptimizer(
      0.1) if optimizer is None else optimizer

  def train_loop_body(step):
    train_op = optimizer.minimize(
        build_loss_fn if tf.executing_eagerly() else build_loss_fn())
    return tf.tuple(tensors=[tf.add(step, 1)], control_inputs=[train_op])

  minimize_op = tf.compat.v1.while_loop(
      cond=lambda step: step < num_steps,
      body=train_loop_body,
      loop_vars=[tf.constant(0)],
      return_same_structure=True)[0]  # Always return a single op.
  return minimize_op

def _minimize_in_graph ( build_loss_fn , num_steps = 200 , optimizer = None ) : optimizer = tf . compat . v1 . train . AdamOptimizer ( 0.1 ) if optimizer is None else optimizer def train_loop_body ( step ) : train_op = optimizer . minimize ( build_loss_fn if tf . executing_eagerly ( ) else build_loss_fn ( ) ) return tf . tuple ( tensors = [ tf . add ( step , 1 ) ] , control_inputs = [ train_op ] ) minimize_op = tf . compat . v1 . while_loop ( cond = lambda step : step < num_steps , body = train_loop_body , loop_vars = [ tf . constant ( 0 ) ] , return_same_structure = True ) [ 0 ] return minimize_op

Run an optimizer within the graph to minimize a loss function .


def moments_of_masked_time_series(time_series_tensor, broadcast_mask):
  """Compute mean and variance, accounting for a mask.

  Args:
    time_series_tensor: float `Tensor` time series of shape
      `concat([batch_shape, [num_timesteps]])`.
    broadcast_mask: bool `Tensor` of the same shape as `time_series`.
  Returns:
    mean: float `Tensor` of shape `batch_shape`.
    variance: float `Tensor` of shape `batch_shape`.
  """
  num_unmasked_entries = tf.cast(
      tf.reduce_sum(input_tensor=tf.cast(~broadcast_mask, tf.int32), axis=-1),
      time_series_tensor.dtype)

  # Manually compute mean and variance, excluding masked entries.
  mean = (tf.reduce_sum(input_tensor=tf.where(
      broadcast_mask,
      tf.zeros_like(time_series_tensor),
      time_series_tensor), axis=-1) / num_unmasked_entries)
  variance = (tf.reduce_sum(input_tensor=tf.where(
      broadcast_mask,
      tf.zeros_like(time_series_tensor),
      (time_series_tensor - mean[..., tf.newaxis]) ** 2), axis=-1)
              / num_unmasked_entries)
  return mean, variance

def moments_of_masked_time_series ( time_series_tensor , broadcast_mask ) : num_unmasked_entries = tf . cast ( tf . reduce_sum ( input_tensor = tf . cast ( ~ broadcast_mask , tf . int32 ) , axis = - 1 ) , time_series_tensor . dtype ) mean = ( tf . reduce_sum ( input_tensor = tf . where ( broadcast_mask , tf . zeros_like ( time_series_tensor ) , time_series_tensor ) , axis = - 1 ) / num_unmasked_entries ) variance = ( tf . reduce_sum ( input_tensor = tf . where ( broadcast_mask , tf . zeros_like ( time_series_tensor ) , ( time_series_tensor - mean [ ... , tf . newaxis ] ) ** 2 ) , axis = - 1 ) / num_unmasked_entries ) return mean , variance

Compute mean and variance accounting for a mask .


def initial_value_of_masked_time_series(time_series_tensor, broadcast_mask):
  """Get the first unmasked entry of each time series in the batch.

  Args:
    time_series_tensor: float `Tensor` of shape [..., num_timesteps].
    broadcast_mask: bool `Tensor` of same shape as `time_series`.
  """

  num_timesteps = tf.shape(input=time_series_tensor)[-1]

  # Compute the index of the first unmasked entry for each series in the batch.
  unmasked_negindices = (
      tf.cast(~broadcast_mask, tf.int32) *
      tf.range(num_timesteps, 0, -1))
  first_unmasked_indices = num_timesteps - tf.reduce_max(
      input_tensor=unmasked_negindices, axis=-1)

  if first_unmasked_indices.shape.ndims is None:
    raise NotImplementedError(
        'Cannot compute initial values of a masked time series with'
        'dynamic rank.')  # `batch_gather` requires static rank

  # Extract the initial value for each series in the batch.
  return tf.squeeze(tf.compat.v1.batch_gather(
      params=time_series_tensor,
      indices=first_unmasked_indices[..., tf.newaxis]), axis=-1)

Get the first unmasked entry of each time series in the batch .


def broadcast_batch_shape(distributions):
  """Get broadcast batch shape from distributions, statically if possible."""

  # Static case
  batch_shape = distributions[0].batch_shape
  for distribution in distributions:
    batch_shape = tf.broadcast_static_shape(batch_shape,
                                            distribution.batch_shape)
  if batch_shape.is_fully_defined():
    return batch_shape.as_list()

  # Fallback on dynamic.
  batch_shape = distributions[0].batch_shape_tensor()
  for distribution in distributions:
    batch_shape = tf.broadcast_dynamic_shape(batch_shape,
                                             distribution.batch_shape_tensor())

  return tf.convert_to_tensor(value=batch_shape)

def broadcast_batch_shape ( distributions ) : batch_shape = distributions [ 0 ] . batch_shape for distribution in distributions : batch_shape = tf . broadcast_static_shape ( batch_shape , distribution . batch_shape ) if batch_shape . is_fully_defined ( ) : return batch_shape . as_list ( ) batch_shape = distributions [ 0 ] . batch_shape_tensor ( ) for distribution in distributions : batch_shape = tf . broadcast_dynamic_shape ( batch_shape , distribution . batch_shape_tensor ( ) ) return tf . convert_to_tensor ( value = batch_shape )

Get broadcast batch shape from distributions statically if possible .


def factored_joint_mvn(distributions):
  """Combine MultivariateNormals into a factored joint distribution.

   Given a list of multivariate normal distributions
   `dist[i] = Normal(loc[i], scale[i])`, construct the joint
   distribution given by concatenating independent samples from these
   distributions. This is multivariate normal with mean vector given by the
   concatenation of the component mean vectors, and block-diagonal covariance
   matrix in which the blocks are the component covariances.

   Note that for computational efficiency, multivariate normals are represented
   by a 'scale' (factored covariance) linear operator rather than the full
   covariance matrix.

  Args:
    distributions: Python `iterable` of MultivariateNormal distribution
      instances (e.g., `tfd.MultivariateNormalDiag`,
      `tfd.MultivariateNormalTriL`, etc.). These must be broadcastable to a
      consistent batch shape, but may have different event shapes
      (i.e., defined over spaces of different dimension).

  Returns:
    joint_distribution: An instance of `tfd.MultivariateNormalLinearOperator`
      representing the joint distribution constructed by concatenating
      an independent sample from each input distributions.
  """

  graph_parents = [tensor for distribution in distributions
                   for tensor in distribution._graph_parents]  # pylint: disable=protected-access
  with tf.compat.v1.name_scope('factored_joint_mvn', values=graph_parents):

    # We explicitly broadcast the `locs` so that we can concatenate them.
    # We don't have direct numerical access to the `scales`, which are arbitrary
    # linear operators, but `LinearOperatorBlockDiag` appears to do the right
    # thing without further intervention.
    dtype = tf.debugging.assert_same_float_dtype(distributions)
    broadcast_ones = tf.ones(broadcast_batch_shape(distributions),
                             dtype=dtype)[..., tf.newaxis]
    return MultivariateNormalLinearOperator(
        loc=tf.concat([mvn.mean() * broadcast_ones for mvn in distributions],
                      axis=-1),
        scale=tfl.LinearOperatorBlockDiag([mvn.scale for mvn in distributions],
                                          is_square=True))

def factored_joint_mvn ( distributions ) : graph_parents = [ tensor for distribution in distributions for tensor in distribution . _graph_parents ] with tf . compat . v1 . name_scope ( 'factored_joint_mvn' , values = graph_parents ) : dtype = tf . debugging . assert_same_float_dtype ( distributions ) broadcast_ones = tf . ones ( broadcast_batch_shape ( distributions ) , dtype = dtype ) [ ... , tf . newaxis ] return MultivariateNormalLinearOperator ( loc = tf . concat ( [ mvn . mean ( ) * broadcast_ones for mvn in distributions ] , axis = - 1 ) , scale = tfl . LinearOperatorBlockDiag ( [ mvn . scale for mvn in distributions ] , is_square = True ) )

Combine MultivariateNormals into a factored joint distribution .


def sum_mvns(distributions):
  """Attempt to sum MultivariateNormal distributions.

  The sum of (multivariate) normal random variables is itself (multivariate)
  normal, with mean given by the sum of means and (co)variance given by the
  sum of (co)variances. This method exploits this fact to compute the
  sum of a list of `tfd.MultivariateNormalDiag` objects.

  It may in the future be extended to support summation of other forms of
  (Multivariate)Normal distributions.

  Args:
    distributions: Python `iterable` of `tfd.MultivariateNormalDiag`
      distribution instances. These must all have the same event
      shape, and broadcast to a consistent batch shape.

  Returns:
    sum_distribution: A `tfd.MultivariateNormalDiag` instance with mean
      equal to the sum of input means and covariance equal to the sum of
      input covariances.
  """

  graph_parents = [tensor for distribution in distributions
                   for tensor in distribution._graph_parents]  # pylint: disable=protected-access
  with tf.compat.v1.name_scope('sum_mvns', values=graph_parents):
    if all([isinstance(mvn, tfd.MultivariateNormalDiag)
            for mvn in distributions]):
      return tfd.MultivariateNormalDiag(
          loc=sum([mvn.mean() for mvn in distributions]),
          scale_diag=tf.sqrt(sum([
              mvn.scale.diag**2 for mvn in distributions])))
    else:
      raise NotImplementedError(
          'Sums of distributions other than MultivariateNormalDiag are not '
          'currently implemented. (given: {})'.format(distributions))

Attempt to sum MultivariateNormal distributions .


def empirical_statistics(observed_time_series):
  """Compute statistics of a provided time series, as heuristic initialization.

  Args:
    observed_time_series: `Tensor` representing a time series, or batch of time
       series, of shape either `batch_shape + [num_timesteps, 1]` or
       `batch_shape + [num_timesteps]` (allowed if `num_timesteps > 1`).

  Returns:
    observed_mean: `Tensor` of shape `batch_shape`, giving the empirical
      mean of each time series in the batch.
    observed_stddev: `Tensor` of shape `batch_shape`, giving the empirical
      standard deviation of each time series in the batch.
    observed_initial_centered: `Tensor of shape `batch_shape`, giving the
      initial value of each time series in the batch after centering
      (subtracting the mean).
  """

  with tf.compat.v1.name_scope(
      'empirical_statistics', values=[observed_time_series]):

    [
        observed_time_series,
        mask
    ] = canonicalize_observed_time_series_with_mask(observed_time_series)

    squeezed_series = observed_time_series[..., 0]
    if mask is None:
      observed_mean, observed_variance = tf.nn.moments(
          x=squeezed_series, axes=-1)
      observed_initial = squeezed_series[..., 0]
    else:
      broadcast_mask = tf.broadcast_to(tf.cast(mask, tf.bool),
                                       tf.shape(input=squeezed_series))
      observed_mean, observed_variance = (
          missing_values_util.moments_of_masked_time_series(
              squeezed_series, broadcast_mask=broadcast_mask))
      try:
        observed_initial = (
            missing_values_util.initial_value_of_masked_time_series(
                squeezed_series, broadcast_mask=broadcast_mask))
      except NotImplementedError:
        tf.compat.v1.logging.warn(
            'Cannot compute initial values for a masked time series'
            'with dynamic shape; using the mean instead. This will'
            'affect heuristic priors and may change the results of'
            'inference.')
        observed_initial = observed_mean

    observed_stddev = tf.sqrt(observed_variance)
    observed_initial_centered = observed_initial - observed_mean
    return observed_mean, observed_stddev, observed_initial_centered

Compute statistics of a provided time series as heuristic initialization .


def _maybe_expand_trailing_dim(observed_time_series_tensor):
  """Ensures `observed_time_series_tensor` has a trailing dimension of size 1.

  The `tfd.LinearGaussianStateSpaceModel` Distribution has event shape of
  `[num_timesteps, observation_size]`, but canonical BSTS models
  are univariate, so their observation_size is always `1`. The extra trailing
  dimension gets annoying, so this method allows arguments with or without the
  extra dimension. There is no ambiguity except in the trivial special case
  where  `num_timesteps = 1`; this can be avoided by specifying any unit-length
  series in the explicit `[num_timesteps, 1]` style.

  Most users should not call this method directly, and instead call
  `canonicalize_observed_time_series_with_mask`, which handles converting
  to `Tensor` and specifying an optional missingness mask.

  Args:
    observed_time_series_tensor: `Tensor` of shape
      `batch_shape + [num_timesteps, 1]` or `batch_shape + [num_timesteps]`,
      where `num_timesteps > 1`.

  Returns:
    expanded_time_series: `Tensor` of shape `batch_shape + [num_timesteps, 1]`.
  """

  with tf.compat.v1.name_scope(
      'maybe_expand_trailing_dim', values=[observed_time_series_tensor]):
    if (observed_time_series_tensor.shape.ndims is not None and
        tf.compat.dimension_value(
            observed_time_series_tensor.shape[-1]) is not None):
      expanded_time_series = (
          observed_time_series_tensor
          if observed_time_series_tensor.shape[-1] == 1
          else observed_time_series_tensor[..., tf.newaxis])
    else:
      expanded_time_series = tf.cond(
          pred=tf.equal(tf.shape(input=observed_time_series_tensor)[-1], 1),
          true_fn=lambda: observed_time_series_tensor,
          false_fn=lambda: observed_time_series_tensor[..., tf.newaxis])
    return expanded_time_series

def _maybe_expand_trailing_dim ( observed_time_series_tensor ) : with tf . compat . v1 . name_scope ( 'maybe_expand_trailing_dim' , values = [ observed_time_series_tensor ] ) : if ( observed_time_series_tensor . shape . ndims is not None and tf . compat . dimension_value ( observed_time_series_tensor . shape [ - 1 ] ) is not None ) : expanded_time_series = ( observed_time_series_tensor if observed_time_series_tensor . shape [ - 1 ] == 1 else observed_time_series_tensor [ ... , tf . newaxis ] ) else : expanded_time_series = tf . cond ( pred = tf . equal ( tf . shape ( input = observed_time_series_tensor ) [ - 1 ] , 1 ) , true_fn = lambda : observed_time_series_tensor , false_fn = lambda : observed_time_series_tensor [ ... , tf . newaxis ] ) return expanded_time_series

Ensures observed_time_series_tensor has a trailing dimension of size 1 .


def canonicalize_observed_time_series_with_mask(
    maybe_masked_observed_time_series):
  """Extract a Tensor with canonical shape and optional mask.

  Args:
    maybe_masked_observed_time_series: a `Tensor`-like object with shape
      `[..., num_timesteps]` or `[..., num_timesteps, 1]`, or a
      `tfp.sts.MaskedTimeSeries` containing such an object.
  Returns:
    masked_time_series: a `tfp.sts.MaskedTimeSeries` namedtuple, in which
      the `observed_time_series` is converted to `Tensor` with canonical shape
      `[..., num_timesteps, 1]`, and `is_missing` is either `None` or a boolean
      `Tensor`.
  """

  with tf.compat.v1.name_scope('canonicalize_observed_time_series_with_mask'):
    if hasattr(maybe_masked_observed_time_series, 'is_missing'):
      observed_time_series = (
          maybe_masked_observed_time_series.time_series)
      is_missing = maybe_masked_observed_time_series.is_missing
    else:
      observed_time_series = maybe_masked_observed_time_series
      is_missing = None

    observed_time_series = tf.convert_to_tensor(value=observed_time_series,
                                                name='observed_time_series')
    observed_time_series = _maybe_expand_trailing_dim(observed_time_series)

    if is_missing is not None:
      is_missing = tf.convert_to_tensor(
          value=is_missing, name='is_missing', dtype_hint=tf.bool)

    return missing_values_util.MaskedTimeSeries(observed_time_series,
                                                is_missing=is_missing)

Extract a Tensor with canonical shape and optional mask .


def mix_over_posterior_draws(means, variances):
  """Construct a predictive normal distribution that mixes over posterior draws.

  Args:
    means: float `Tensor` of shape
      `[num_posterior_draws, ..., num_timesteps]`.
    variances: float `Tensor` of shape
      `[num_posterior_draws, ..., num_timesteps]`.

  Returns:
    mixture_dist: `tfd.MixtureSameFamily(tfd.Independent(tfd.Normal))` instance
      representing a uniform mixture over the posterior samples, with
      `batch_shape = ...` and `event_shape = [num_timesteps]`.

  """
  # The inputs `means`, `variances` have shape
  #   `concat([
  #      [num_posterior_draws],
  #      sample_shape,
  #      batch_shape,
  #      [num_timesteps]])`
  # Because MixtureSameFamily mixes over the rightmost batch dimension,
  # we need to move the `num_posterior_draws` dimension to be rightmost
  # in the batch shape. This requires use of `Independent` (to preserve
  # `num_timesteps` as part of the event shape) and `move_dimension`.
  # TODO(b/120245392): enhance `MixtureSameFamily` to reduce along an
  # arbitrary axis, and eliminate `move_dimension` calls here.

  with tf.compat.v1.name_scope(
      'mix_over_posterior_draws', values=[means, variances]):
    num_posterior_draws = dist_util.prefer_static_value(
        tf.shape(input=means))[0]

    component_observations = tfd.Independent(
        distribution=tfd.Normal(
            loc=dist_util.move_dimension(means, 0, -2),
            scale=tf.sqrt(dist_util.move_dimension(variances, 0, -2))),
        reinterpreted_batch_ndims=1)

    return tfd.MixtureSameFamily(
        mixture_distribution=tfd.Categorical(
            logits=tf.zeros([num_posterior_draws],
                            dtype=component_observations.dtype)),
        components_distribution=component_observations)

def mix_over_posterior_draws ( means , variances ) : with tf . compat . v1 . name_scope ( 'mix_over_posterior_draws' , values = [ means , variances ] ) : num_posterior_draws = dist_util . prefer_static_value ( tf . shape ( input = means ) ) [ 0 ] component_observations = tfd . Independent ( distribution = tfd . Normal ( loc = dist_util . move_dimension ( means , 0 , - 2 ) , scale = tf . sqrt ( dist_util . move_dimension ( variances , 0 , - 2 ) ) ) , reinterpreted_batch_ndims = 1 ) return tfd . MixtureSameFamily ( mixture_distribution = tfd . Categorical ( logits = tf . zeros ( [ num_posterior_draws ] , dtype = component_observations . dtype ) ) , components_distribution = component_observations )

Construct a predictive normal distribution that mixes over posterior draws .


def range(self, name="range"):
    """`high - low`."""
    with self._name_scope(name):
      return self.high - self.low

def range ( self , name = "range" ) : with self . _name_scope ( name ) : return self . high - self . low


def _make_summary_statistic(attr):
  """Factory for making summary statistics, eg, mean, mode, stddev."""
  def _fn(self):
    if any(self._dist_fn_args):  # pylint: disable=protected-access
      raise ValueError(
          'Can only compute ' + attr + ' when all distributions are '
          'independent; {}'.format(self.model))
    return self._unflatten(getattr(d(), attr)() for d in self._dist_fn_wrapped)  # pylint: disable=protected-access
  return _fn

def
_make_summary_statistic
(
attr
)
:
def
_fn
(
self
)
:
if
any
(
self
.
_dist_fn_args
)
:
raise
ValueError
(
'Can only compute '
+
attr
+
' when all distributions are '
'independent; {}'
.
format
(
self
.
model
)
)
return
self
.
_unflatten
(
getattr
(
d
(
)
,
attr
)
(
)
for
d
in
self
.
_dist_fn_wrapped
)
return
_fn

Factory for making summary statistics eg mean mode stddev .


def _unify_call_signature(i, dist_fn):
  """Creates `dist_fn_wrapped` which calls `dist_fn` with all prev nodes.

  Args:
    i: Python `int` corresponding to position in topologically sorted DAG.
    dist_fn: Python `callable` which takes a subset of previously constructed
      distributions (in reverse order) and produces a new distribution instance.

  Returns:
    dist_fn_wrapped: Python `callable` which takes all previous distributions
      (in non reverse order) and produces a  new distribution instance.
    args: `tuple` of `str` representing the arg names of `dist_fn` (and in non
      wrapped, "natural" order). `None` is returned only if the input is not a
      `callable`.
  """
  if distribution_util.is_distribution_instance(dist_fn):
    return (lambda *_: dist_fn), None

  if not callable(dist_fn):
    raise TypeError('{} must be either `tfd.Distribution`-like or '
                    '`callable`.'.format(dist_fn))

  args = _get_required_args(dist_fn)
  if not args:
    return (lambda *_: dist_fn()), ()

  @functools.wraps(dist_fn)
  def dist_fn_wrapped(*xs):
    """Calls `dist_fn` with reversed and truncated args."""
    if i != len(xs):
      raise ValueError(
          'Internal Error: Unexpected number of inputs provided to {}-th '
          'distribution maker (dist_fn: {}, expected: {}, saw: {}).'.format(
              i, dist_fn, i, len(xs)))
    if len(xs) < len(args):
      raise ValueError(
          'Internal Error: Too few inputs provided to {}-th distribution maker '
          '(dist_fn: {}, expected: {}, saw: {}).'.format(
              i, dist_fn, len(args), len(xs)))
    return dist_fn(*reversed(xs[-len(args):]))
  return dist_fn_wrapped, args

Creates dist_fn_wrapped which calls dist_fn with all prev nodes .


def _resolve_distribution_names(dist_fn_args, dist_names, leaf_name):
  """Uses arg names to resolve distribution names."""
  if dist_names is None:
    dist_names = []
  else:
    dist_names = dist_names.copy()
  n = len(dist_fn_args)
  dist_names.extend([None]*(n - len(dist_names)))
  for i_, args in enumerate(reversed(dist_fn_args)):
    if not args:
      continue  # There's no args to analyze.
    i = n - i_ - 1
    for j, arg_name in enumerate(args):
      dist_names[i - j - 1] = arg_name
  j = 0
  for i_ in range(len(dist_names)):
    i = n - i_ - 1
    if dist_names[i] is None:
      dist_names[i] = leaf_name if j == 0 else leaf_name + str(j)
      j += 1
  return tuple(dist_names)

def _resolve_distribution_names ( dist_fn_args , dist_names , leaf_name ) : if dist_names is None : dist_names = [ ] else : dist_names = dist_names . copy ( ) n = len ( dist_fn_args ) dist_names . extend ( [ None ] * ( n - len ( dist_names ) ) ) for i_ , args in enumerate ( reversed ( dist_fn_args ) ) : if not args : continue i = n - i_ - 1 for j , arg_name in enumerate ( args ) : dist_names [ i - j - 1 ] = arg_name j = 0 for i_ in range ( len ( dist_names ) ) : i = n - i_ - 1 if dist_names [ i ] is None : dist_names [ i ] = leaf_name if j == 0 else leaf_name + str ( j ) j += 1 return tuple ( dist_names )

Uses arg names to resolve distribution names .


def _get_required_args(fn):
  """Returns the distribution's required args."""
  argspec = tf_inspect.getfullargspec(fn)
  args = argspec.args
  if tf_inspect.isclass(fn):
    args = args[1:]  # Remove the `self` arg.
  if argspec.defaults:
    # Remove the args which have defaults. By convention we only feed
    # *required args*. This means some distributions must always be wrapped
    # with a `lambda`, e.g., `lambda logits: tfd.Bernoulli(logits=logits)`
    # or `lambda probs: tfd.Bernoulli(probs=probs)`.
    args = args[:-len(argspec.defaults)]
  return tuple(args)

def _get_required_args ( fn ) : argspec = tf_inspect . getfullargspec ( fn ) args = argspec . args if tf_inspect . isclass ( fn ) : args = args [ 1 : ] if argspec . defaults : args = args [ : - len ( argspec . defaults ) ] return tuple ( args )

Returns the distribution s required args .


def _kl_joint_joint(d0, d1, name=None):
  """Calculate the KL divergence between two `JointDistributionSequential`s.

  Args:
    d0: instance of a `JointDistributionSequential` object.
    d1: instance of a `JointDistributionSequential` object.
    name: (optional) Name to use for created operations.
      Default value: `"kl_joint_joint"`.

  Returns:
    kl_joint_joint: `Tensor` The sum of KL divergences between elemental
      distributions of two joint distributions.

  Raises:
    ValueError: when joint distributions have a different number of elemental
      distributions.
    ValueError: when either joint distribution has a distribution with dynamic
      dependency, i.e., when either joint distribution is not a collection of
      independent distributions.
  """
  if len(d0._dist_fn_wrapped) != len(d1._dist_fn_wrapped):  # pylint: disable=protected-access
    raise ValueError(
        'Can only compute KL divergence between when each has the'
        'same number of component distributions.')
  if (not all(a is None for a in d0._dist_fn_args) or  # pylint: disable=protected-access
      not all(a is None for a in d1._dist_fn_args)):  # pylint: disable=protected-access
    raise ValueError(
        'Can only compute KL divergence when all distributions are '
        'independent.')
  with tf.name_scope(name or 'kl_jointseq_jointseq'):
    return sum(kullback_leibler.kl_divergence(d0_(), d1_())
               for d0_, d1_ in zip(d0._dist_fn_wrapped, d1._dist_fn_wrapped))

Calculate the KL divergence between two JointDistributionSequential s .


def _build(self, model):
    """Creates `dist_fn`, `dist_fn_wrapped`, `dist_fn_args`."""
    if not isinstance(model, collections.Sequence):
      raise TypeError('`model` must be `list`-like (saw: {}).'.format(
          type(model).__name__))
    self._dist_fn = model
    self._dist_fn_wrapped, self._dist_fn_args = zip(*[
        _unify_call_signature(i, dist_fn)
        for i, dist_fn in enumerate(model)])


def _resolve_graph(self, distribution_names=None, leaf_name='x'):
    """Creates a `tuple` of `tuple`s of dependencies.

    This function is **experimental**. That said, we encourage its use
    and ask that you report problems to `tfprobability@tensorflow.org`.

    Args:
      distribution_names: `list` of `str` or `None` names corresponding to each
        of `model` elements. (`None`s are expanding into the
        appropriate `str`.)
      leaf_name: `str` used when no maker depends on a particular
        `model` element.

    Returns:
      graph: `tuple` of `(str tuple)` pairs representing the name of each
        distribution (maker) and the names of its dependencies.

    #### Example

    ```python
    d = tfd.JointDistributionSequential([
                     tfd.Independent(tfd.Exponential(rate=[100, 120]), 1),
        lambda    e: tfd.Gamma(concentration=e[..., 0], rate=e[..., 1]),
                     tfd.Normal(loc=0, scale=2.),
        lambda n, g: tfd.Normal(loc=n, scale=g),
    ])
    d._resolve_graph()
    # ==> (
    #       ('e', ()),
    #       ('g', ('e',)),
    #       ('n', ()),
    #       ('x', ('n', 'g')),
    #     )
    ```

    """
    # This function additionally depends on:
    #   self._dist_fn_args
    #   self._dist_fn_wrapped
    # TODO(b/129008220): Robustify this procedure. Eg, handle collisions better,
    # ignore args prefixed with `_`.
    if distribution_names is None or any(self._dist_fn_args):
      distribution_names = _resolve_distribution_names(
          self._dist_fn_args, distribution_names, leaf_name)
    if len(set(distribution_names)) != len(distribution_names):
      raise ValueError('Distribution names must be unique: {}'.format(
          distribution_names))
    if len(distribution_names) != len(self._dist_fn_wrapped):
      raise ValueError('Distribution names must be 1:1 with `rvs`.')
    return tuple(zip(distribution_names,
                     tuple(() if a is None else a for a in self._dist_fn_args)))

Creates a tuple of tuple s of dependencies .


def _entropy(self):
    """Shannon entropy in nats."""
    if any(self._dist_fn_args):
      raise ValueError(
          'Can only compute entropy when all distributions are independent.')
    return sum(joint_distribution_lib.maybe_check_wont_broadcast(
        (d().entropy() for d in self._dist_fn_wrapped),
        self.validate_args))


def check_arg_in_support(f):
  """Decorator function for argument bounds checking.

  This decorator is meant to be used with methods that require the first
  argument to be in the support of the distribution. If `validate_args` is
  `True`, the method is wrapped with an assertion that the first argument is
  greater than or equal to `loc`, since the support of the half-Cauchy
  distribution is given by `[loc, infinity)`.

  Args:
    f: method to be decorated.

  Returns:
    Returns a decorated method that, when `validate_args` attribute of the class
    is `True`, will assert that all elements in the first argument are within
    the support of the distribution before executing the original method.
  """
  @functools.wraps(f)
  def _check_arg_and_apply_f(*args, **kwargs):
    dist = args[0]
    x = args[1]
    with tf.control_dependencies([
        assert_util.assert_greater_equal(
            x, dist.loc, message="x is not in the support of the distribution")
    ] if dist.validate_args else []):
      return f(*args, **kwargs)
  return _check_arg_and_apply_f

Decorator function for argument bounds checking .


def image_summary(seqs, name, num=None):
  """Visualizes sequences as TensorBoard summaries.

  Args:
    seqs: A tensor of shape [n, t, h, w, c].
    name: String name of this summary.
    num: Integer for the number of examples to visualize. Defaults to
      all examples.
  """
  seqs = tf.clip_by_value(seqs, 0., 1.)
  seqs = tf.unstack(seqs[:num])
  joined_seqs = [tf.concat(tf.unstack(seq), 1) for seq in seqs]
  joined_seqs = tf.expand_dims(tf.concat(joined_seqs, 0), 0)
  tf.compat.v2.summary.image(
      name,
      joined_seqs,
      max_outputs=1,
      step=tf.compat.v1.train.get_or_create_global_step())

def image_summary ( seqs , name , num = None ) : seqs = tf . clip_by_value ( seqs , 0. , 1. ) seqs = tf . unstack ( seqs [ : num ] ) joined_seqs = [ tf . concat ( tf . unstack ( seq ) , 1 ) for seq in seqs ] joined_seqs = tf . expand_dims ( tf . concat ( joined_seqs , 0 ) , 0 ) tf . compat . v2 . summary . image ( name , joined_seqs , max_outputs = 1 , step = tf . compat . v1 . train . get_or_create_global_step ( ) )

Visualizes sequences as TensorBoard summaries .


def visualize_reconstruction(inputs, reconstruct, num=3, name="reconstruction"):
  """Visualizes the reconstruction of inputs in TensorBoard.

  Args:
    inputs: A tensor of the original inputs, of shape [batch, timesteps,
      h, w, c].
    reconstruct: A tensor of a reconstruction of inputs, of shape
      [batch, timesteps, h, w, c].
    num: Integer for the number of examples to visualize.
    name: String name of this summary.
  """
  reconstruct = tf.clip_by_value(reconstruct, 0., 1.)
  inputs_and_reconstruct = tf.concat((inputs[:num], reconstruct[:num]), axis=0)
  image_summary(inputs_and_reconstruct, name)

def visualize_reconstruction ( inputs , reconstruct , num = 3 , name = "reconstruction" ) : reconstruct = tf . clip_by_value ( reconstruct , 0. , 1. ) inputs_and_reconstruct = tf . concat ( ( inputs [ : num ] , reconstruct [ : num ] ) , axis = 0 ) image_summary ( inputs_and_reconstruct , name )

Visualizes the reconstruction of inputs in TensorBoard .


def visualize_qualitative_analysis(inputs, model, samples=1, batch_size=3,
                                   length=8):
  """Visualizes a qualitative analysis of a given model.

  Args:
    inputs: A tensor of the original inputs, of shape [batch, timesteps,
      h, w, c].
    model: A DisentangledSequentialVAE model.
    samples: Number of samples to draw from the latent distributions.
    batch_size: Number of sequences to generate.
    length: Number of timesteps to generate for each sequence.
  """
  average = lambda dist: tf.reduce_mean(
      input_tensor=dist.mean(), axis=0)  # avg over samples
  with tf.compat.v1.name_scope("val_reconstruction"):
    reconstruct = functools.partial(model.reconstruct, inputs=inputs,
                                    samples=samples)
    visualize_reconstruction(inputs, average(reconstruct()))
    visualize_reconstruction(inputs, average(reconstruct(sample_static=True)),
                             name="static_prior")
    visualize_reconstruction(inputs, average(reconstruct(sample_dynamic=True)),
                             name="dynamic_prior")
    visualize_reconstruction(inputs, average(reconstruct(swap_static=True)),
                             name="swap_static")
    visualize_reconstruction(inputs, average(reconstruct(swap_dynamic=True)),
                             name="swap_dynamic")

  with tf.compat.v1.name_scope("generation"):
    generate = functools.partial(model.generate, batch_size=batch_size,
                                 length=length, samples=samples)
    image_summary(average(generate(fix_static=True)), "fix_static")
    image_summary(average(generate(fix_dynamic=True)), "fix_dynamic")

def visualize_qualitative_analysis ( inputs , model , samples = 1 , batch_size = 3 , length = 8 ) : average = lambda dist : tf . reduce_mean ( input_tensor = dist . mean ( ) , axis = 0 ) with tf . compat . v1 . name_scope ( "val_reconstruction" ) : reconstruct = functools . partial ( model . reconstruct , inputs = inputs , samples = samples ) visualize_reconstruction ( inputs , average ( reconstruct ( ) ) ) visualize_reconstruction ( inputs , average ( reconstruct ( sample_static = True ) ) , name = "static_prior" ) visualize_reconstruction ( inputs , average ( reconstruct ( sample_dynamic = True ) ) , name = "dynamic_prior" ) visualize_reconstruction ( inputs , average ( reconstruct ( swap_static = True ) ) , name = "swap_static" ) visualize_reconstruction ( inputs , average ( reconstruct ( swap_dynamic = True ) ) , name = "swap_dynamic" ) with tf . compat . v1 . name_scope ( "generation" ) : generate = functools . partial ( model . generate , batch_size = batch_size , length = length , samples = samples ) image_summary ( average ( generate ( fix_static = True ) ) , "fix_static" ) image_summary ( average ( generate ( fix_dynamic = True ) ) , "fix_dynamic" )

Visualizes a qualitative analysis of a given model .


def summarize_dist_params(dist, name, name_scope="dist_params"):
  """Summarize the parameters of a distribution.

  Args:
    dist: A Distribution object with mean and standard deviation
      parameters.
    name: The name of the distribution.
    name_scope: The name scope of this summary.
  """
  with tf.compat.v1.name_scope(name_scope):
    tf.compat.v2.summary.histogram(
        name="{}/{}".format(name, "mean"),
        data=dist.mean(),
        step=tf.compat.v1.train.get_or_create_global_step())
    tf.compat.v2.summary.histogram(
        name="{}/{}".format(name, "stddev"),
        data=dist.stddev(),
        step=tf.compat.v1.train.get_or_create_global_step())

def summarize_dist_params ( dist , name , name_scope = "dist_params" ) : with tf . compat . v1 . name_scope ( name_scope ) : tf . compat . v2 . summary . histogram ( name = "{}/{}" . format ( name , "mean" ) , data = dist . mean ( ) , step = tf . compat . v1 . train . get_or_create_global_step ( ) ) tf . compat . v2 . summary . histogram ( name = "{}/{}" . format ( name , "stddev" ) , data = dist . stddev ( ) , step = tf . compat . v1 . train . get_or_create_global_step ( ) )

Summarize the parameters of a distribution .


def summarize_mean_in_nats_and_bits(inputs, units, name,
                                    nats_name_scope="nats",
                                    bits_name_scope="bits_per_dim"):
  """Summarize the mean of a tensor in nats and bits per unit.

  Args:
    inputs: A tensor of values measured in nats.
    units: The units of the tensor with which to compute the mean bits
      per unit.
    name: The name of the tensor.
    nats_name_scope: The name scope of the nats summary.
    bits_name_scope: The name scope of the bits summary.
  """
  mean = tf.reduce_mean(input_tensor=inputs)
  with tf.compat.v1.name_scope(nats_name_scope):
    tf.compat.v2.summary.scalar(
        name,
        mean,
        step=tf.compat.v1.train.get_or_create_global_step())
  with tf.compat.v1.name_scope(bits_name_scope):
    tf.compat.v2.summary.scalar(
        name,
        mean / units / tf.math.log(2.),
        step=tf.compat.v1.train.get_or_create_global_step())

def summarize_mean_in_nats_and_bits ( inputs , units , name , nats_name_scope = "nats" , bits_name_scope = "bits_per_dim" ) : mean = tf . reduce_mean ( input_tensor = inputs ) with tf . compat . v1 . name_scope ( nats_name_scope ) : tf . compat . v2 . summary . scalar ( name , mean , step = tf . compat . v1 . train . get_or_create_global_step ( ) ) with tf . compat . v1 . name_scope ( bits_name_scope ) : tf . compat . v2 . summary . scalar ( name , mean / units / tf . math . log ( 2. ) , step = tf . compat . v1 . train . get_or_create_global_step ( ) )

Summarize the mean of a tensor in nats and bits per unit .


def call(self, inputs):
    """Runs the model to generate multivariate normal distribution.

    Args:
      inputs: Unused.

    Returns:
      A MultivariateNormalDiag distribution with event shape
      [dimensions], batch shape [], and sample shape [sample_shape,
      dimensions].
    """
    del inputs  # unused
    with tf.compat.v1.name_scope(self._name):
      return tfd.MultivariateNormalDiag(self.loc, self.scale_diag)

def call ( self , inputs ) : del inputs with tf . compat . v1 . name_scope ( self . _name ) : return tfd . MultivariateNormalDiag ( self . loc , self . scale_diag )

Runs the model to generate multivariate normal distribution .


def zero_state(self, sample_batch_shape=()):
    """Returns an initial state for the LSTM cell.

    Args:
      sample_batch_shape: A 0D or 1D tensor of the combined sample and
        batch shape.

    Returns:
      A tuple of the initial previous output at timestep 0 of shape
      [sample_batch_shape, dimensions], and the cell state.
    """
    h0 = tf.zeros([1, self.hidden_size])
    c0 = tf.zeros([1, self.hidden_size])
    combined_shape = tf.concat((tf.convert_to_tensor(
        value=sample_batch_shape, dtype=tf.int32), [self.dimensions]),
                               axis=-1)
    previous_output = tf.zeros(combined_shape)
    return previous_output, (h0, c0)

def zero_state ( self , sample_batch_shape = ( ) ) : h0 = tf . zeros ( [ 1 , self . hidden_size ] ) c0 = tf . zeros ( [ 1 , self . hidden_size ] ) combined_shape = tf . concat ( ( tf . convert_to_tensor ( value = sample_batch_shape , dtype = tf . int32 ) , [ self . dimensions ] ) , axis = - 1 ) previous_output = tf . zeros ( combined_shape ) return previous_output , ( h0 , c0 )

Returns an initial state for the LSTM cell .


def call(self, inputs, state):
    """Runs the model to generate a distribution for a single timestep.

    This generates a batched MultivariateNormalDiag distribution using
    the output of the recurrent model at the current timestep to
    parameterize the distribution.

    Args:
      inputs: The sampled value of `z` at the previous timestep, i.e.,
        `z_{t-1}`, of shape [..., dimensions].
        `z_0` should be set to the empty matrix.
      state: A tuple containing the (hidden, cell) state.

    Returns:
      A tuple of a MultivariateNormalDiag distribution, and the state of
      the recurrent function at the end of the current timestep. The
      distribution will have event shape [dimensions], batch shape
      [...], and sample shape [sample_shape, ..., dimensions].
    """
    # In order to allow the user to pass in a single example without a batch
    # dimension, we always expand the input to at least two dimensions, then
    # fix the output shape to remove the batch dimension if necessary.
    original_shape = inputs.shape
    if len(original_shape) < 2:
      inputs = tf.reshape(inputs, [1, -1])
    out, state = self.lstm_cell(inputs, state)
    out = self.output_layer(out)
    correct_shape = tf.concat((original_shape[:-1], tf.shape(input=out)[-1:]),
                              0)
    out = tf.reshape(out, correct_shape)
    loc = out[..., :self.dimensions]
    scale_diag = tf.nn.softplus(out[..., self.dimensions:]) + 1e-5  # keep > 0
    return tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale_diag), state

def call ( self , inputs , state ) : original_shape = inputs . shape if len ( original_shape ) < 2 : inputs = tf . reshape ( inputs , [ 1 , - 1 ] ) out , state = self . lstm_cell ( inputs , state ) out = self . output_layer ( out ) correct_shape = tf . concat ( ( original_shape [ : - 1 ] , tf . shape ( input = out ) [ - 1 : ] ) , 0 ) out = tf . reshape ( out , correct_shape ) loc = out [ ... , : self . dimensions ] scale_diag = tf . nn . softplus ( out [ ... , self . dimensions : ] ) + 1e-5 return tfd . MultivariateNormalDiag ( loc = loc , scale_diag = scale_diag ) , state

Runs the model to generate a distribution for a single timestep .


def call(self, inputs):
    """Runs the model to generate an intermediate representation of x_t.

    Args:
      inputs: A batch of image sequences `x_{1:T}` of shape
        `[sample_shape, batch_size, timesteps, height, width,
        channels]`.

    Returns:
      A batch of intermediate representations of shape [sample_shape,
      batch_size, timesteps, hidden_size].
    """
    image_shape = tf.shape(input=inputs)[-3:]
    collapsed_shape = tf.concat(([-1], image_shape), axis=0)
    out = tf.reshape(inputs, collapsed_shape)  # (sample*batch*T, h, w, c)
    out = self.conv1(out)
    out = self.conv2(out)
    out = self.conv3(out)
    out = self.conv4(out)
    expanded_shape = tf.concat((tf.shape(input=inputs)[:-3], [-1]), axis=0)
    return tf.reshape(out, expanded_shape)

def call ( self , inputs ) : image_shape = tf . shape ( input = inputs ) [ - 3 : ] collapsed_shape = tf . concat ( ( [ - 1 ] , image_shape ) , axis = 0 ) out = tf . reshape ( inputs , collapsed_shape ) out = self . conv1 ( out ) out = self . conv2 ( out ) out = self . conv3 ( out ) out = self . conv4 ( out ) expanded_shape = tf . concat ( ( tf . shape ( input = inputs ) [ : - 3 ] , [ - 1 ] ) , axis = 0 ) return tf . reshape ( out , expanded_shape )

Runs the model to generate an intermediate representation of x_t .


def generate(self, batch_size, length, samples=1, fix_static=False,
               fix_dynamic=False):
    """Generate new sequences.

    Args:
      batch_size: Number of sequences to generate.
      length: Number of timesteps to generate for each sequence.
      samples: Number of samples to draw from the latent distributions.
      fix_static: Boolean for whether or not to share the same random
        sample of the static latent variable `f` from its prior across
        all examples.
      fix_dynamic: Boolean for whether or not to share the same random
        sample of the dynamic latent variable `z_{1:T}` from its prior
        across all examples.

    Returns:
      A batched Independent distribution wrapping a set of Normal
      distributions over the pixels of the generated sequences, where
      the Independent distribution has event shape [height, width,
      channels], batch shape [samples, batch_size, timesteps], and
      sample shape [sample_shape, samples, batch_size, timesteps,
      height, width, channels].
    """
    static_sample, _ = self.sample_static_prior(samples, batch_size, fix_static)
    dynamic_sample, _ = self.sample_dynamic_prior(samples, batch_size, length,
                                                  fix_dynamic)
    likelihood = self.decoder((dynamic_sample, static_sample))
    return likelihood

def generate ( self , batch_size , length , samples = 1 , fix_static = False , fix_dynamic = False ) : static_sample , _ = self . sample_static_prior ( samples , batch_size , fix_static ) dynamic_sample , _ = self . sample_dynamic_prior ( samples , batch_size , length , fix_dynamic ) likelihood = self . decoder ( ( dynamic_sample , static_sample ) ) return likelihood


def reconstruct(self, inputs, samples=1, sample_static=False,
                  sample_dynamic=False, swap_static=False, swap_dynamic=False,
                  fix_static=False, fix_dynamic=False):
    """Reconstruct the given input sequences.

    Args:
      inputs: A batch of image sequences `x_{1:T}` of shape
        `[batch_size, timesteps, height, width, channels]`.
      samples: Number of samples to draw from the latent distributions.
      sample_static: Boolean for whether or not to randomly sample the
        static latent variable `f` from its prior distribution.
      sample_dynamic: Boolean for whether or not to randomly sample the
        dynamic latent variable `z_{1:T}` from its prior distribution.
      swap_static: Boolean for whether or not to swap the encodings for
        the static latent variable `f` between the examples.
      swap_dynamic: Boolean for whether or not to swap the encodings for
        the dynamic latent variable `z_{1:T}` between the examples.
      fix_static: Boolean for whether or not to share the same random
        sample of the static latent variable `f` from its prior across
        all examples.
      fix_dynamic: Boolean for whether or not to share the same random
        sample of the dynamic latent variable `z_{1:T}` from its prior
        across all examples.

    Returns:
      A batched Independent distribution wrapping a set of Normal
      distributions over the pixels of the reconstruction of the input,
      where the Independent distribution has event shape [height, width,
      channels], batch shape [samples, batch_size, timesteps], and
      sample shape [sample_shape, samples, batch_size, timesteps,
      height, width, channels].
    """
    batch_size = tf.shape(input=inputs)[-5]
    length = len(tf.unstack(inputs, axis=-4))  # hack for graph mode

    features = self.compressor(inputs)  # (..., batch, timesteps, hidden)

    if sample_static:
      static_sample, _ = self.sample_static_prior(
          samples, batch_size, fix_static)
    else:
      static_sample, _ = self.sample_static_posterior(features, samples)

    if swap_static:
      static_sample = tf.reverse(static_sample, axis=[1])

    if sample_dynamic:
      dynamic_sample, _ = self.sample_dynamic_prior(
          samples, batch_size, length, fix_dynamic)
    else:
      dynamic_sample, _ = self.sample_dynamic_posterior(
          features, samples, static_sample)

    if swap_dynamic:
      dynamic_sample = tf.reverse(dynamic_sample, axis=[1])

    likelihood = self.decoder((dynamic_sample, static_sample))
    return likelihood

def reconstruct ( self , inputs , samples = 1 , sample_static = False , sample_dynamic = False , swap_static = False , swap_dynamic = False , fix_static = False , fix_dynamic = False ) : batch_size = tf . shape ( input = inputs ) [ - 5 ] length = len ( tf . unstack ( inputs , axis = - 4 ) ) features = self . compressor ( inputs ) if sample_static : static_sample , _ = self . sample_static_prior ( samples , batch_size , fix_static ) else : static_sample , _ = self . sample_static_posterior ( features , samples ) if swap_static : static_sample = tf . reverse ( static_sample , axis = [ 1 ] ) if sample_dynamic : dynamic_sample , _ = self . sample_dynamic_prior ( samples , batch_size , length , fix_dynamic ) else : dynamic_sample , _ = self . sample_dynamic_posterior ( features , samples , static_sample ) if swap_dynamic : dynamic_sample = tf . reverse ( dynamic_sample , axis = [ 1 ] ) likelihood = self . decoder ( ( dynamic_sample , static_sample ) ) return likelihood


def sample_static_prior(self, samples, batch_size, fixed=False):
    """Sample the static latent prior.

    Args:
      samples: Number of samples to draw from the latent distribution.
      batch_size: Number of sequences to sample.
      fixed: Boolean for whether or not to share the same random
        sample across all sequences.

    Returns:
      A tuple of a sample tensor of shape [samples, batch_size,
      latent_size], and a MultivariateNormalDiag distribution from which
      the tensor was sampled, with event shape [latent_size], and batch
      shape [].
    """
    dist = self.static_prior()
    if fixed:  # in either case, shape is (samples, batch, latent)
      sample = dist.sample((samples, 1)) + tf.zeros([batch_size, 1])
    else:
      sample = dist.sample((samples, batch_size))
    return sample, dist

def sample_static_prior ( self , samples , batch_size , fixed = False ) : dist = self . static_prior ( ) if fixed : sample = dist . sample ( ( samples , 1 ) ) + tf . zeros ( [ batch_size , 1 ] ) else : sample = dist . sample ( ( samples , batch_size ) ) return sample , dist


def sample_dynamic_prior(self, samples, batch_size, length, fixed=False):
    """Sample the dynamic latent prior.

    Args:
      samples: Number of samples to draw from the latent distribution.
      batch_size: Number of sequences to sample.
      length: Number of timesteps to sample for each sequence.
      fixed: Boolean for whether or not to share the same random
        sample across all sequences.

    Returns:
      A tuple of a sample tensor of shape [samples, batch_size, length
      latent_size], and a MultivariateNormalDiag distribution from which
      the tensor was sampled, with event shape [latent_size], and batch
      shape [samples, 1, length] if fixed or [samples, batch_size,
      length] otherwise.
    """
    if fixed:
      sample_batch_size = 1
    else:
      sample_batch_size = batch_size

    sample, state = self.dynamic_prior.zero_state([samples, sample_batch_size])
    locs = []
    scale_diags = []
    sample_list = []
    for _ in range(length):
      dist, state = self.dynamic_prior(sample, state)
      sample = dist.sample()
      locs.append(dist.parameters["loc"])
      scale_diags.append(dist.parameters["scale_diag"])
      sample_list.append(sample)

    sample = tf.stack(sample_list, axis=2)
    loc = tf.stack(locs, axis=2)
    scale_diag = tf.stack(scale_diags, axis=2)

    if fixed:  # tile along the batch axis
      sample = sample + tf.zeros([batch_size, 1, 1])

    return sample, tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale_diag)

def sample_dynamic_prior ( self , samples , batch_size , length , fixed = False ) : if fixed : sample_batch_size = 1 else : sample_batch_size = batch_size sample , state = self . dynamic_prior . zero_state ( [ samples , sample_batch_size ] ) locs = [ ] scale_diags = [ ] sample_list = [ ] for _ in range ( length ) : dist , state = self . dynamic_prior ( sample , state ) sample = dist . sample ( ) locs . append ( dist . parameters [ "loc" ] ) scale_diags . append ( dist . parameters [ "scale_diag" ] ) sample_list . append ( sample ) sample = tf . stack ( sample_list , axis = 2 ) loc = tf . stack ( locs , axis = 2 ) scale_diag = tf . stack ( scale_diags , axis = 2 ) if fixed : sample = sample + tf . zeros ( [ batch_size , 1 , 1 ] ) return sample , tfd . MultivariateNormalDiag ( loc = loc , scale_diag = scale_diag )


def batch_shape(self):
    """Static batch shape of models represented by this component.

    Returns:
      batch_shape: A `tf.TensorShape` giving the broadcast batch shape of
        all model parameters. This should match the batch shape of
        derived state space models, i.e.,
        `self.make_state_space_model(...).batch_shape`. It may be partially
        defined or unknown.
    """
    batch_shape = tf.TensorShape([])
    for param in self.parameters:
      batch_shape = tf.broadcast_static_shape(
          batch_shape, param.prior.batch_shape)
    return batch_shape

def batch_shape ( self ) : batch_shape = tf . TensorShape ( [ ] ) for param in self . parameters : batch_shape = tf . broadcast_static_shape ( batch_shape , param . prior . batch_shape ) return batch_shape

Static batch shape of models represented by this component .


def batch_shape_tensor(self):
    """Runtime batch shape of models represented by this component.

    Returns:
      batch_shape: `int` `Tensor` giving the broadcast batch shape of
        all model parameters. This should match the batch shape of
        derived state space models, i.e.,
        `self.make_state_space_model(...).batch_shape_tensor()`.
    """
    batch_shape = tf.constant([], dtype=tf.int32)
    for param in self.parameters:
      batch_shape = tf.broadcast_dynamic_shape(
          batch_shape, param.prior.batch_shape_tensor())
    return batch_shape

def batch_shape_tensor ( self ) : batch_shape = tf . constant ( [ ] , dtype = tf . int32 ) for param in self . parameters : batch_shape = tf . broadcast_dynamic_shape ( batch_shape , param . prior . batch_shape_tensor ( ) ) return batch_shape

Runtime batch shape of models represented by this component .


def make_state_space_model(self,
                             num_timesteps,
                             param_vals=None,
                             initial_state_prior=None,
                             initial_step=0):
    """Instantiate this model as a Distribution over specified `num_timesteps`.

    Args:
      num_timesteps: Python `int` number of timesteps to model.
      param_vals: a list of `Tensor` parameter values in order corresponding to
        `self.parameters`, or a dict mapping from parameter names to values.
      initial_state_prior: an optional `Distribution` instance overriding the
        default prior on the model's initial state. This is used in forecasting
        ("today's prior is yesterday's posterior").
      initial_step: optional `int` specifying the initial timestep to model.
        This is relevant when the model contains time-varying components,
        e.g., holidays or seasonality.

    Returns:
      dist: a `LinearGaussianStateSpaceModel` Distribution object.
    """
    return self._make_state_space_model(
        num_timesteps=num_timesteps,
        param_map=self._canonicalize_param_vals_as_map(param_vals),
        initial_state_prior=initial_state_prior,
        initial_step=initial_step)

Instantiate this model as a Distribution over specified num_timesteps .


def prior_sample(self,
                   num_timesteps,
                   initial_step=0,
                   params_sample_shape=(),
                   trajectories_sample_shape=(),
                   seed=None):
    """Sample from the joint prior over model parameters and trajectories.

    Args:
      num_timesteps: Scalar `int` `Tensor` number of timesteps to model.
      initial_step: Optional scalar `int` `Tensor` specifying the starting
        timestep.
          Default value: 0.
      params_sample_shape: Number of possible worlds to sample iid from the
        parameter prior, or more generally, `Tensor` `int` shape to fill with
        iid samples.
          Default value: [] (i.e., draw a single sample and don't expand the
          shape).
      trajectories_sample_shape: For each sampled set of parameters, number
        of trajectories to sample, or more generally, `Tensor` `int` shape to
        fill with iid samples.
        Default value: [] (i.e., draw a single sample and don't expand the
          shape).
      seed: Python `int` random seed.

    Returns:
      trajectories: `float` `Tensor` of shape
        `trajectories_sample_shape + params_sample_shape + [num_timesteps, 1]`
        containing all sampled trajectories.
      param_samples: list of sampled parameter value `Tensor`s, in order
        corresponding to `self.parameters`, each of shape
        `params_sample_shape + prior.batch_shape + prior.event_shape`.
    """

    seed = distributions.SeedStream(
        seed, salt='StructuralTimeSeries_prior_sample')

    with tf.compat.v1.name_scope(
        'prior_sample',
        values=[num_timesteps, params_sample_shape, trajectories_sample_shape]):
      param_samples = [
          p.prior.sample(params_sample_shape, seed=seed(), name=p.name)
          for p in self.parameters
      ]
      model = self.make_state_space_model(
          num_timesteps=num_timesteps,
          initial_step=initial_step,
          param_vals=param_samples)
      return model.sample(trajectories_sample_shape, seed=seed()), param_samples

def prior_sample ( self , num_timesteps , initial_step = 0 , params_sample_shape = ( ) , trajectories_sample_shape = ( ) , seed = None ) : seed = distributions . SeedStream ( seed , salt = 'StructuralTimeSeries_prior_sample' ) with tf . compat . v1 . name_scope ( 'prior_sample' , values = [ num_timesteps , params_sample_shape , trajectories_sample_shape ] ) : param_samples = [ p . prior . sample ( params_sample_shape , seed = seed ( ) , name = p . name ) for p in self . parameters ] model = self . make_state_space_model ( num_timesteps = num_timesteps , initial_step = initial_step , param_vals = param_samples ) return model . sample ( trajectories_sample_shape , seed = seed ( ) ) , param_samples

Sample from the joint prior over model parameters and trajectories .


def _compute_min_event_ndims(bijector_list, compute_forward=True):
  """Computes the min_event_ndims associated with the give list of bijectors.

  Given a list `bijector_list` of bijectors, compute the min_event_ndims that is
  associated with the composition of bijectors in that list.

  min_event_ndims is the # of right most dimensions for which the bijector has
  done necessary computation on (i.e. the non-broadcastable part of the
  computation).

  We can derive the min_event_ndims for a chain of bijectors as follows:

  In the case where there are no rank changing bijectors, this will simply be
  `max(b.forward_min_event_ndims for b in bijector_list)`. This is because the
  bijector with the most forward_min_event_ndims requires the most dimensions,
  and hence the chain also requires operating on those dimensions.

  However in the case of rank changing, more care is needed in determining the
  exact amount of dimensions. Padding dimensions causes subsequent bijectors to
  operate on the padded dimensions, and Removing dimensions causes bijectors to
  operate more left.

  Args:
    bijector_list: List of bijectors to be composed by chain.
    compute_forward: Boolean. If True, computes the min_event_ndims associated
      with a forward call to Chain, and otherwise computes the min_event_ndims
      associated with an inverse call to Chain. The latter is the same as the
      min_event_ndims associated with a forward call to Invert(Chain(....)).

  Returns:
    min_event_ndims
  """
  min_event_ndims = 0
  # This is a mouthful, but what this encapsulates is that if not for rank
  # changing bijectors, we'd only need to compute the largest of the min
  # required ndims. Hence "max_min". Due to rank changing bijectors, we need to
  # account for synthetic rank growth / synthetic rank decrease from a rank
  # changing bijector.
  rank_changed_adjusted_max_min_event_ndims = 0

  if compute_forward:
    bijector_list = reversed(bijector_list)

  for b in bijector_list:
    if compute_forward:
      current_min_event_ndims = b.forward_min_event_ndims
      current_inverse_min_event_ndims = b.inverse_min_event_ndims
    else:
      current_min_event_ndims = b.inverse_min_event_ndims
      current_inverse_min_event_ndims = b.forward_min_event_ndims

    # New dimensions were touched.
    if rank_changed_adjusted_max_min_event_ndims < current_min_event_ndims:
      min_event_ndims += (
          current_min_event_ndims - rank_changed_adjusted_max_min_event_ndims)
    rank_changed_adjusted_max_min_event_ndims = max(
        current_min_event_ndims, rank_changed_adjusted_max_min_event_ndims)

    # If the number of dimensions has increased via forward, then
    # inverse_min_event_ndims > forward_min_event_ndims, and hence the
    # dimensions we computed on, have moved left (so we have operated
    # on additional dimensions).
    # Conversely, if the number of dimensions has decreased via forward,
    # then we have inverse_min_event_ndims < forward_min_event_ndims,
    # and so we will have operated on fewer right most dimensions.

    number_of_changed_dimensions = (
        current_min_event_ndims - current_inverse_min_event_ndims)
    rank_changed_adjusted_max_min_event_ndims -= number_of_changed_dimensions
  return min_event_ndims

Computes the min_event_ndims associated with the give list of bijectors .


def vector_size_to_square_matrix_size(d, validate_args, name=None):
  """Convert a vector size to a matrix size."""
  if isinstance(d, (float, int, np.generic, np.ndarray)):
    n = (-1 + np.sqrt(1 + 8 * d)) / 2.
    if float(int(n)) != n:
      raise ValueError("Vector length is not a triangular number.")
    return int(n)
  else:
    with tf.name_scope(name or "vector_size_to_square_matrix_size") as name:
      n = (-1. + tf.sqrt(1 + 8. * tf.cast(d, dtype=tf.float32))) / 2.
      if validate_args:
        with tf.control_dependencies([
            assert_util.assert_equal(
                tf.cast(tf.cast(n, dtype=tf.int32), dtype=tf.float32),
                n,
                message="Vector length is not a triangular number")
        ]):
          n = tf.identity(n)
      return tf.cast(n, d.dtype)

Convert a vector size to a matrix size .


def _argsort(values, axis=-1, direction='ASCENDING', stable=False, name=None):  # pylint: disable=unused-argument
  """Numpy implementation of `tf.argsort`."""
  if direction == 'ASCENDING':
    pass
  elif direction == 'DESCENDING':
    values = np.negative(values)
  else:
    raise ValueError('Unrecognized direction: {}.'.format(direction))
  return np.argsort(values, axis, kind='stable' if stable else 'quicksort')

def
_argsort
(
values
,
axis
=
1
,
direction
=
'ASCENDING'
,
stable
=
False
,
name
=
)
:
if
direction
==
'ASCENDING'
:
pass
elif
direction
==
'DESCENDING'
:
values
=
np
.
negative
(
values
)
else
:
raise
ValueError
(
'Unrecognized direction: {}.'
.
format
(
direction
)
)
return
np
.
argsort
(
values
,
axis
,
kind
=
'stable'
if
stable
else
'quicksort'
)


def _sort(values, axis=-1, direction='ASCENDING', stable=False, name=None):  # pylint: disable=unused-argument
  """Numpy implementation of `tf.sort`."""
  if direction == 'ASCENDING':
    pass
  elif direction == 'DESCENDING':
    values = np.negative(values)
  else:
    raise ValueError('Unrecognized direction: {}.'.format(direction))
  result = np.sort(values, axis, kind='stable' if stable else 'quicksort')
  if direction == 'DESCENDING':
    return np.negative(result)
  return result

def
_sort
(
values
,
axis
=
1
,
direction
=
'ASCENDING'
,
stable
=
False
,
name
=
)
:
if
direction
==
'ASCENDING'
:
pass
elif
direction
==
'DESCENDING'
:
values
=
np
.
negative
(
values
)
else
:
raise
ValueError
(
'Unrecognized direction: {}.'
.
format
(
direction
)
)
result
=
np
.
sort
(
values
,
axis
,
kind
=
'stable'
if
stable
else
'quicksort'
)
if
direction
==
'DESCENDING'
:
return
np
.
negative
(
result
)
return
result


def ndtr(x, name="ndtr"):
  """Normal distribution function.

  Returns the area under the Gaussian probability density function, integrated
  from minus infinity to x:

  ```
                    1       / x
     ndtr(x)  = ----------  |    exp(-0.5 t**2) dt
                sqrt(2 pi)  /-inf

              = 0.5 (1 + erf(x / sqrt(2)))
              = 0.5 erfc(x / sqrt(2))
  ```

  Args:
    x: `Tensor` of type `float32`, `float64`.
    name: Python string. A name for the operation (default="ndtr").

  Returns:
    ndtr: `Tensor` with `dtype=x.dtype`.

  Raises:
    TypeError: if `x` is not floating-type.
  """

  with tf.name_scope(name):
    x = tf.convert_to_tensor(value=x, name="x")
    if dtype_util.as_numpy_dtype(x.dtype) not in [np.float32, np.float64]:
      raise TypeError(
          "x.dtype=%s is not handled, see docstring for supported types."
          % x.dtype)
    return _ndtr(x)


def _ndtr(x):
  """Implements ndtr core logic."""
  half_sqrt_2 = tf.constant(
      0.5 * np.sqrt(2.), dtype=x.dtype, name="half_sqrt_2")
  w = x * half_sqrt_2
  z = tf.abs(w)
  y = tf.where(
      tf.less(z, half_sqrt_2), 1. + tf.math.erf(w),
      tf.where(tf.greater(w, 0.), 2. - tf.math.erfc(z), tf.math.erfc(z)))
  return 0.5 * y

def _ndtr ( x ) : half_sqrt_2 = tf . constant ( 0.5 * np . sqrt ( 2. ) , dtype = x . dtype , name = "half_sqrt_2" ) w = x * half_sqrt_2 z = tf . abs ( w ) y = tf . where ( tf . less ( z , half_sqrt_2 ) , 1. + tf . math . erf ( w ) , tf . where ( tf . greater ( w , 0. ) , 2. - tf . math . erfc ( z ) , tf . math . erfc ( z ) ) ) return 0.5 * y


def ndtri(p, name="ndtri"):
  """The inverse of the CDF of the Normal distribution function.

  Returns x such that the area under the pdf from minus infinity to x is equal
  to p.

  A piece-wise rational approximation is done for the function.
  This is a port of the implementation in netlib.

  Args:
    p: `Tensor` of type `float32`, `float64`.
    name: Python string. A name for the operation (default="ndtri").

  Returns:
    x: `Tensor` with `dtype=p.dtype`.

  Raises:
    TypeError: if `p` is not floating-type.
  """

  with tf.name_scope(name):
    p = tf.convert_to_tensor(value=p, name="p")
    if dtype_util.as_numpy_dtype(p.dtype) not in [np.float32, np.float64]:
      raise TypeError(
          "p.dtype=%s is not handled, see docstring for supported types."
          % p.dtype)
    return _ndtri(p)

The inverse of the CDF of the Normal distribution function .


def log_ndtr(x, series_order=3, name="log_ndtr"):
  """Log Normal distribution function.

  For details of the Normal distribution function see `ndtr`.

  This function calculates `(log o ndtr)(x)` by either calling `log(ndtr(x))` or
  using an asymptotic series. Specifically:
  - For `x > upper_segment`, use the approximation `-ndtr(-x)` based on
    `log(1-x) ~= -x, x << 1`.
  - For `lower_segment < x <= upper_segment`, use the existing `ndtr` technique
    and take a log.
  - For `x <= lower_segment`, we use the series approximation of erf to compute
    the log CDF directly.

  The `lower_segment` is set based on the precision of the input:

  ```
  lower_segment = { -20,  x.dtype=float64
                  { -10,  x.dtype=float32
  upper_segment = {   8,  x.dtype=float64
                  {   5,  x.dtype=float32
  ```

  When `x < lower_segment`, the `ndtr` asymptotic series approximation is:

  ```
     ndtr(x) = scale * (1 + sum) + R_N
     scale   = exp(-0.5 x**2) / (-x sqrt(2 pi))
     sum     = Sum{(-1)^n (2n-1)!! / (x**2)^n, n=1:N}
     R_N     = O(exp(-0.5 x**2) (2N+1)!! / |x|^{2N+3})
  ```

  where `(2n-1)!! = (2n-1) (2n-3) (2n-5) ...  (3) (1)` is a
  [double-factorial](https://en.wikipedia.org/wiki/Double_factorial).

  Args:
    x: `Tensor` of type `float32`, `float64`.
    series_order: Positive Python `integer`. Maximum depth to
      evaluate the asymptotic expansion. This is the `N` above.
    name: Python string. A name for the operation (default="log_ndtr").

  Returns:
    log_ndtr: `Tensor` with `dtype=x.dtype`.

  Raises:
    TypeError: if `x.dtype` is not handled.
    TypeError: if `series_order` is a not Python `integer.`
    ValueError:  if `series_order` is not in `[0, 30]`.
  """
  if not isinstance(series_order, int):
    raise TypeError("series_order must be a Python integer.")
  if series_order < 0:
    raise ValueError("series_order must be non-negative.")
  if series_order > 30:
    raise ValueError("series_order must be <= 30.")

  with tf.name_scope(name):
    x = tf.convert_to_tensor(value=x, name="x")

    if dtype_util.base_equal(x.dtype, tf.float64):
      lower_segment = LOGNDTR_FLOAT64_LOWER
      upper_segment = LOGNDTR_FLOAT64_UPPER
    elif dtype_util.base_equal(x.dtype, tf.float32):
      lower_segment = LOGNDTR_FLOAT32_LOWER
      upper_segment = LOGNDTR_FLOAT32_UPPER
    else:
      raise TypeError("x.dtype=%s is not supported." % x.dtype)

    # The basic idea here was ported from:
    #   https://root.cern.ch/doc/v608/SpecFuncCephesInv_8cxx_source.html
    # We copy the main idea, with a few changes
    # * For x >> 1, and X ~ Normal(0, 1),
    #     Log[P[X < x]] = Log[1 - P[X < -x]] approx -P[X < -x],
    #     which extends the range of validity of this function.
    # * We use one fixed series_order for all of 'x', rather than adaptive.
    # * Our docstring properly reflects that this is an asymptotic series, not a
    #   Taylor series. We also provided a correct bound on the remainder.
    # * We need to use the max/min in the _log_ndtr_lower arg to avoid nan when
    #   x=0. This happens even though the branch is unchosen because when x=0
    #   the gradient of a select involves the calculation 1*dy+0*(-inf)=nan
    #   regardless of whether dy is finite. Note that the minimum is a NOP if
    #   the branch is chosen.
    return tf.where(
        tf.greater(x, upper_segment),
        -_ndtr(-x),  # log(1-x) ~= -x, x << 1
        tf.where(
            tf.greater(x, lower_segment),
            tf.math.log(_ndtr(tf.maximum(x, lower_segment))),
            _log_ndtr_lower(tf.minimum(x, lower_segment), series_order)))


def _log_ndtr_asymptotic_series(x, series_order):
  """Calculates the asymptotic series used in log_ndtr."""
  npdt = dtype_util.as_numpy_dtype(x.dtype)
  if series_order <= 0:
    return npdt(1)
  x_2 = tf.square(x)
  even_sum = tf.zeros_like(x)
  odd_sum = tf.zeros_like(x)
  x_2n = x_2  # Start with x^{2*1} = x^{2*n} with n = 1.
  for n in range(1, series_order + 1):
    y = npdt(_double_factorial(2 * n - 1)) / x_2n
    if n % 2:
      odd_sum += y
    else:
      even_sum += y
    x_2n *= x_2
  return 1. + even_sum - odd_sum

def _log_ndtr_asymptotic_series ( x , series_order ) : npdt = dtype_util . as_numpy_dtype ( x . dtype ) if series_order <= 0 : return npdt ( 1 ) x_2 = tf . square ( x ) even_sum = tf . zeros_like ( x ) odd_sum = tf . zeros_like ( x ) x_2n = x_2 for n in range ( 1 , series_order + 1 ) : y = npdt ( _double_factorial ( 2 * n - 1 ) ) / x_2n if n % 2 : odd_sum += y else : even_sum += y x_2n *= x_2 return 1. + even_sum - odd_sum

Calculates the asymptotic series used in log_ndtr .


def erfinv(x, name="erfinv"):
  """The inverse function for erf, the error function.

  Args:
    x: `Tensor` of type `float32`, `float64`.
    name: Python string. A name for the operation (default="erfinv").

  Returns:
    x: `Tensor` with `dtype=x.dtype`.

  Raises:
    TypeError: if `x` is not floating-type.
  """

  with tf.name_scope(name):
    x = tf.convert_to_tensor(value=x, name="x")
    if dtype_util.as_numpy_dtype(x.dtype) not in [np.float32, np.float64]:
      raise TypeError("x.dtype={} is not handled, see docstring for supported "
                      "types.".format(dtype_util.name(x.dtype)))
    return ndtri((x + 1.) / 2.) / np.sqrt(2.)

The inverse function for erf the error function .


def log_cdf_laplace(x, name="log_cdf_laplace"):
  """Log Laplace distribution function.

  This function calculates `Log[L(x)]`, where `L(x)` is the cumulative
  distribution function of the Laplace distribution, i.e.

  ```L(x) := 0.5 * int_{-infty}^x e^{-|t|} dt```

  For numerical accuracy, `L(x)` is computed in different ways depending on `x`,

  ```
  x <= 0:
    Log[L(x)] = Log[0.5] + x, which is exact

  0 < x:
    Log[L(x)] = Log[1 - 0.5 * e^{-x}], which is exact
  ```

  Args:
    x: `Tensor` of type `float32`, `float64`.
    name: Python string. A name for the operation (default="log_ndtr").

  Returns:
    `Tensor` with `dtype=x.dtype`.

  Raises:
    TypeError: if `x.dtype` is not handled.
  """

  with tf.name_scope(name):
    x = tf.convert_to_tensor(value=x, name="x")

    # For x < 0, L(x) = 0.5 * exp{x} exactly, so Log[L(x)] = log(0.5) + x.
    lower_solution = -np.log(2.) + x

    # safe_exp_neg_x = exp{-x} for x > 0, but is
    # bounded above by 1, which avoids
    #   log[1 - 1] = -inf for x = log(1/2), AND
    #   exp{-x} --> inf, for x << -1
    safe_exp_neg_x = tf.exp(-tf.abs(x))

    # log1p(z) = log(1 + z) approx z for |z| << 1. This approxmation is used
    # internally by log1p, rather than being done explicitly here.
    upper_solution = tf.math.log1p(-0.5 * safe_exp_neg_x)

    return tf.where(x < 0., lower_solution, upper_solution)

def log_cdf_laplace ( x , name = "log_cdf_laplace" ) : with tf . name_scope ( name ) : x = tf . convert_to_tensor ( value = x , name = "x" ) lower_solution = - np . log ( 2. ) + x safe_exp_neg_x = tf . exp ( - tf . abs ( x ) ) upper_solution = tf . math . log1p ( - 0.5 * safe_exp_neg_x ) return tf . where ( x < 0. , lower_solution , upper_solution )


def text_messages_joint_log_prob(count_data, lambda_1, lambda_2, tau):
  """Joint log probability function."""
  alpha = (1. / tf.reduce_mean(input_tensor=count_data))
  rv_lambda = tfd.Exponential(rate=alpha)

  rv_tau = tfd.Uniform()

  lambda_ = tf.gather(
      [lambda_1, lambda_2],
      indices=tf.cast(
          tau * tf.cast(tf.size(input=count_data), dtype=tf.float32) <= tf.cast(
              tf.range(tf.size(input=count_data)), dtype=tf.float32),
          dtype=tf.int32))
  rv_observation = tfd.Poisson(rate=lambda_)

  return (rv_lambda.log_prob(lambda_1) + rv_lambda.log_prob(lambda_2) +
          rv_tau.log_prob(tau) +
          tf.reduce_sum(input_tensor=rv_observation.log_prob(count_data)))

def text_messages_joint_log_prob ( count_data , lambda_1 , lambda_2 , tau ) : alpha = ( 1. / tf . reduce_mean ( input_tensor = count_data ) ) rv_lambda = tfd . Exponential ( rate = alpha ) rv_tau = tfd . Uniform ( ) lambda_ = tf . gather ( [ lambda_1 , lambda_2 ] , indices = tf . cast ( tau * tf . cast ( tf . size ( input = count_data ) , dtype = tf . float32 ) <= tf . cast ( tf . range ( tf . size ( input = count_data ) ) , dtype = tf . float32 ) , dtype = tf . int32 ) ) rv_observation = tfd . Poisson ( rate = lambda_ ) return ( rv_lambda . log_prob ( lambda_1 ) + rv_lambda . log_prob ( lambda_2 ) + rv_tau . log_prob ( tau ) + tf . reduce_sum ( input_tensor = rv_observation . log_prob ( count_data ) ) )


def benchmark_text_messages_hmc(
    num_results=int(3e3),
    num_burnin_steps=int(3e3),
    num_leapfrog_steps=3):
  """Runs HMC on the text-messages unnormalized posterior."""

  if not tf.executing_eagerly():
    tf.compat.v1.reset_default_graph()

  # Build a static, pretend dataset.
  count_data = tf.cast(
      tf.concat(
          [tfd.Poisson(rate=15.).sample(43),
           tfd.Poisson(rate=25.).sample(31)],
          axis=0),
      dtype=tf.float32)
  if tf.executing_eagerly():
    count_data = count_data.numpy()
  else:
    with tf.compat.v1.Session():
      count_data = count_data.eval()

  # Define a closure over our joint_log_prob.
  def unnormalized_log_posterior(lambda1, lambda2, tau):
    return text_messages_joint_log_prob(count_data, lambda1, lambda2, tau)

  if tf.executing_eagerly():
    sample_chain = tf.function(tfp.mcmc.sample_chain)
  else:
    sample_chain = tfp.mcmc.sample_chain

  # Initialize the step_size. (It will be automatically adapted.)
  step_size = tf.compat.v2.Variable(
      name='step_size',
      initial_value=tf.constant(0.05, dtype=tf.float32),
      trainable=False)

  def computation():
    """The benchmark computation."""

    initial_chain_state = [
        tf.constant(count_data.mean(), name='init_lambda1'),
        tf.constant(count_data.mean(), name='init_lambda2'),
        tf.constant(0.5, name='init_tau'),
    ]

    unconstraining_bijectors = [
        tfp.bijectors.Exp(),       # Maps a positive real to R.
        tfp.bijectors.Exp(),       # Maps a positive real to R.
        tfp.bijectors.Sigmoid(),   # Maps [0,1] to R.
    ]

    _, kernel_results = sample_chain(
        num_results=num_results,
        num_burnin_steps=num_burnin_steps,
        current_state=initial_chain_state,
        kernel=tfp.mcmc.TransformedTransitionKernel(
            inner_kernel=tfp.mcmc.HamiltonianMonteCarlo(
                target_log_prob_fn=unnormalized_log_posterior,
                num_leapfrog_steps=num_leapfrog_steps,
                step_size=step_size,
                step_size_update_fn=
                tfp.mcmc.make_simple_step_size_update_policy(num_burnin_steps),
                state_gradients_are_stopped=True),
            bijector=unconstraining_bijectors))

    return kernel_results.inner_results.is_accepted

  # Let's force evaluation of graph to ensure build time is not part of our time
  # trial.
  is_accepted_tensor = computation()
  if not tf.executing_eagerly():
    session = tf.compat.v1.Session()
    session.run(tf.compat.v1.global_variables_initializer())
    session.run(is_accepted_tensor)

  start_time = time.time()
  if tf.executing_eagerly():
    is_accepted = computation()
  else:
    is_accepted = session.run(is_accepted_tensor)
  wall_time = time.time() - start_time

  num_accepted = np.sum(is_accepted)
  acceptance_rate = np.float32(num_accepted) / np.float32(num_results)

  return dict(
      iters=(num_results + num_burnin_steps) * num_leapfrog_steps,
      extras={'acceptance_rate': acceptance_rate},
      wall_time=wall_time)

def benchmark_text_messages_hmc ( num_results = int ( 3e3 ) , num_burnin_steps = int ( 3e3 ) , num_leapfrog_steps = 3 ) : if not tf . executing_eagerly ( ) : tf . compat . v1 . reset_default_graph ( ) count_data = tf . cast ( tf . concat ( [ tfd . Poisson ( rate = 15. ) . sample ( 43 ) , tfd . Poisson ( rate = 25. ) . sample ( 31 ) ] , axis = 0 ) , dtype = tf . float32 ) if tf . executing_eagerly ( ) : count_data = count_data . numpy ( ) else : with tf . compat . v1 . Session ( ) : count_data = count_data . eval ( ) def unnormalized_log_posterior ( lambda1 , lambda2 , tau ) : return text_messages_joint_log_prob ( count_data , lambda1 , lambda2 , tau ) if tf . executing_eagerly ( ) : sample_chain = tf . function ( tfp . mcmc . sample_chain ) else : sample_chain = tfp . mcmc . sample_chain step_size = tf . compat . v2 . Variable ( name = 'step_size' , initial_value = tf . constant ( 0.05 , dtype = tf . float32 ) , trainable = False ) def computation ( ) : initial_chain_state = [ tf . constant ( count_data . mean ( ) , name = 'init_lambda1' ) , tf . constant ( count_data . mean ( ) , name = 'init_lambda2' ) , tf . constant ( 0.5 , name = 'init_tau' ) , ] unconstraining_bijectors = [ tfp . bijectors . Exp ( ) , tfp . bijectors . Exp ( ) , tfp . bijectors . Sigmoid ( ) , ] _ , kernel_results = sample_chain ( num_results = num_results , num_burnin_steps = num_burnin_steps , current_state = initial_chain_state , kernel = tfp . mcmc . TransformedTransitionKernel ( inner_kernel = tfp . mcmc . HamiltonianMonteCarlo ( target_log_prob_fn = unnormalized_log_posterior , num_leapfrog_steps = num_leapfrog_steps , step_size = step_size , step_size_update_fn = tfp . mcmc . make_simple_step_size_update_policy ( num_burnin_steps ) , state_gradients_are_stopped = True ) , bijector = unconstraining_bijectors ) ) return kernel_results . inner_results . is_accepted is_accepted_tensor = computation ( ) if not tf . executing_eagerly ( ) : session = tf . compat . v1 . Session ( ) session . run ( tf . compat . v1 . global_variables_initializer ( ) ) session . run ( is_accepted_tensor ) start_time = time . time ( ) if tf . executing_eagerly ( ) : is_accepted = computation ( ) else : is_accepted = session . run ( is_accepted_tensor ) wall_time = time . time ( ) - start_time num_accepted = np . sum ( is_accepted ) acceptance_rate = np . float32 ( num_accepted ) / np . float32 ( num_results ) return dict ( iters = ( num_results + num_burnin_steps ) * num_leapfrog_steps , extras = { 'acceptance_rate' : acceptance_rate } , wall_time = wall_time )

Runs HMC on the text - messages unnormalized posterior .


def _is_univariate_marginal(self, index_points):
    """True if the given index_points would yield a univariate marginal.

    Args:
      index_points: the set of index set locations at which to compute the
      marginal Gaussian distribution. If this set is of size 1, the marginal is
      univariate.

    Returns:
      is_univariate: Boolean indicating whether the marginal is univariate or
      multivariate. In the case of dynamic shape in the number of index points,
      defaults to "multivariate" since that's the best we can do.
    """
    num_index_points = tf.compat.dimension_value(
        index_points.shape[-(self.kernel.feature_ndims + 1)])
    if num_index_points is None:
      warnings.warn(
          'Unable to detect statically whether the number of index_points is '
          '1. As a result, defaulting to treating the marginal GP at '
          '`index_points` as a multivariate Gaussian. This makes some methods, '
          'like `cdf` unavailable.')
    return num_index_points == 1

True if the given index_points would yield a univariate marginal .


def get_marginal_distribution(self, index_points=None):
    """Compute the marginal of this GP over function values at `index_points`.

    Args:
      index_points: `float` `Tensor` representing finite (batch of) vector(s) of
        points in the index set over which the GP is defined. Shape has the form
        `[b1, ..., bB, e, f1, ..., fF]` where `F` is the number of feature
        dimensions and must equal `kernel.feature_ndims` and `e` is the number
        (size) of index points in each batch. Ultimately this distribution
        corresponds to a `e`-dimensional multivariate normal. The batch shape
        must be broadcastable with `kernel.batch_shape` and any batch dims
        yielded by `mean_fn`.

    Returns:
      marginal: a `Normal` or `MultivariateNormalLinearOperator` distribution,
        according to whether `index_points` consists of one or many index
        points, respectively.
    """
    with self._name_scope('get_marginal_distribution'):
      # TODO(cgs): consider caching the result here, keyed on `index_points`.
      index_points = self._get_index_points(index_points)
      covariance = self._compute_covariance(index_points)
      loc = self._mean_fn(index_points)
      # If we're sure the number of index points is 1, we can just construct a
      # scalar Normal. This has computational benefits and supports things like
      # CDF that aren't otherwise straightforward to provide.
      if self._is_univariate_marginal(index_points):
        scale = tf.sqrt(covariance)
        # `loc` has a trailing 1 in the shape; squeeze it.
        loc = tf.squeeze(loc, axis=-1)
        return normal.Normal(
            loc=loc,
            scale=scale,
            validate_args=self._validate_args,
            allow_nan_stats=self._allow_nan_stats,
            name='marginal_distribution')
      else:
        scale = tf.linalg.LinearOperatorLowerTriangular(
            tf.linalg.cholesky(_add_diagonal_shift(covariance, self.jitter)),
            is_non_singular=True,
            name='GaussianProcessScaleLinearOperator')
        return mvn_linear_operator.MultivariateNormalLinearOperator(
            loc=loc,
            scale=scale,
            validate_args=self._validate_args,
            allow_nan_stats=self._allow_nan_stats,
            name='marginal_distribution')

def get_marginal_distribution ( self , index_points = None ) : with self . _name_scope ( 'get_marginal_distribution' ) : index_points = self . _get_index_points ( index_points ) covariance = self . _compute_covariance ( index_points ) loc = self . _mean_fn ( index_points ) if self . _is_univariate_marginal ( index_points ) : scale = tf . sqrt ( covariance ) loc = tf . squeeze ( loc , axis = - 1 ) return normal . Normal ( loc = loc , scale = scale , validate_args = self . _validate_args , allow_nan_stats = self . _allow_nan_stats , name = 'marginal_distribution' ) else : scale = tf . linalg . LinearOperatorLowerTriangular ( tf . linalg . cholesky ( _add_diagonal_shift ( covariance , self . jitter ) ) , is_non_singular = True , name = 'GaussianProcessScaleLinearOperator' ) return mvn_linear_operator . MultivariateNormalLinearOperator ( loc = loc , scale = scale , validate_args = self . _validate_args , allow_nan_stats = self . _allow_nan_stats , name = 'marginal_distribution' )

Compute the marginal of this GP over function values at index_points .


def _get_index_points(self, index_points=None):
    """Return `index_points` if not None, else `self._index_points`.

    Args:
      index_points: if given, this is what is returned; else,
      `self._index_points`

    Returns:
      index_points: the given arg, if not None, else the class member
      `self._index_points`.

    Rases:
      ValueError: if `index_points` and `self._index_points` are both `None`.
    """
    if self._index_points is None and index_points is None:
      raise ValueError(
          'This GaussianProcess instance was not instantiated with a value for '
          'index_points. One must therefore be provided when calling sample, '
          'log_prob, and other such methods. In particular, one can\'t compute '
          'KL divergences to/from an instance of `GaussianProccess` with '
          'unspecified `index_points` directly. Instead, use the '
          '`get_marginal_distribution` function, which takes `index_points` as '
          'an argument and returns a `Normal` or '
          '`MultivariateNormalLinearOperator` instance, whose KL can be '
          'computed.')
    return index_points if index_points is not None else self._index_points

Return index_points if not None else self . _index_points .


def make_iaf_stack(total_event_size,
                   num_hidden_layers=2,
                   seed=None,
                   dtype=tf.float32):
  """Creates an stacked IAF bijector.

  This bijector operates on vector-valued events.

  Args:
    total_event_size: Number of dimensions to operate over.
    num_hidden_layers: How many hidden layers to use in each IAF.
    seed: Random seed for the initializers.
    dtype: DType for the variables.

  Returns:
    bijector: The created bijector.
  """

  seed = tfd.SeedStream(seed, 'make_iaf_stack')

  def make_iaf():
    """Create an IAF."""
    initializer = tf.compat.v2.keras.initializers.VarianceScaling(
        2 * 0.01, seed=seed() % (2**31 - 1))

    made = tfb.AutoregressiveLayer(
        params=2,
        event_shape=[total_event_size],
        hidden_units=[total_event_size] * num_hidden_layers,
        activation=tf.nn.elu,
        kernel_initializer=initializer,
        dtype=dtype)

    def shift_and_scale(x):
      # TODO(siege): Something is losing the static shape.
      x.set_shape(
          x.shape.merge_with([None] * (x.shape.ndims - 1) + [total_event_size]))
      return tf.unstack(made(x), num=2, axis=-1)

    return tfb.Invert(tfb.MaskedAutoregressiveFlow(shift_and_scale))

  def make_swap():
    """Create an swap."""
    permutation = list(reversed(range(total_event_size)))
    return tfb.Permute(permutation)

  bijector = make_iaf()
  bijector = make_swap()(bijector)
  bijector = make_iaf()(bijector)
  bijector = make_swap()(bijector)
  bijector = make_iaf()(bijector)
  bijector = make_swap()(bijector)

  return bijector

def make_iaf_stack ( total_event_size , num_hidden_layers = 2 , seed = None , dtype = tf . float32 ) : seed = tfd . SeedStream ( seed , 'make_iaf_stack' ) def make_iaf ( ) : initializer = tf . compat . v2 . keras . initializers . VarianceScaling ( 2 * 0.01 , seed = seed ( ) % ( 2 ** 31 - 1 ) ) made = tfb . AutoregressiveLayer ( params = 2 , event_shape = [ total_event_size ] , hidden_units = [ total_event_size ] * num_hidden_layers , activation = tf . nn . elu , kernel_initializer = initializer , dtype = dtype ) def shift_and_scale ( x ) : x . set_shape ( x . shape . merge_with ( [ None ] * ( x . shape . ndims - 1 ) + [ total_event_size ] ) ) return tf . unstack ( made ( x ) , num = 2 , axis = - 1 ) return tfb . Invert ( tfb . MaskedAutoregressiveFlow ( shift_and_scale ) ) def make_swap ( ) : permutation = list ( reversed ( range ( total_event_size ) ) ) return tfb . Permute ( permutation ) bijector = make_iaf ( ) bijector = make_swap ( ) ( bijector ) bijector = make_iaf ( ) ( bijector ) bijector = make_swap ( ) ( bijector ) bijector = make_iaf ( ) ( bijector ) bijector = make_swap ( ) ( bijector ) return bijector


def one_step(self, current_state, previous_kernel_results):
    """Runs one iteration of NeuTra.

    Args:
      current_state: `Tensor` or Python `list` of `Tensor`s representing the
        current state(s) of the Markov chain(s). The first `r` dimensions index
        independent chains, `r = tf.rank(target_log_prob_fn(*current_state))`.
      previous_kernel_results: `collections.namedtuple` containing `Tensor`s
        representing values from previous calls to this function (or from the
        `bootstrap_results` function.)

    Returns:
      next_state: Tensor or Python list of `Tensor`s representing the state(s)
        of the Markov chain(s) after taking exactly one step. Has same type and
        shape as `current_state`.
      kernel_results: `collections.namedtuple` of internal calculations used to
        advance the chain.
    """

    @tfp.mcmc.internal.util.make_innermost_setter
    def set_num_leapfrog_steps(kernel_results, num_leapfrog_steps):
      return kernel_results._replace(
          accepted_results=kernel_results.accepted_results._replace(
              num_leapfrog_steps=num_leapfrog_steps))

    step_size = previous_kernel_results.new_step_size
    previous_kernel_results = set_num_leapfrog_steps(
        previous_kernel_results, self._num_leapfrog_steps(step_size))

    new_state, kernel_results = self._kernel.one_step(
        self._flatten_state(current_state), previous_kernel_results)
    return self._unflatten_state(new_state), kernel_results