with tf.compat.v1.name_scope(<EOL>'<STR_LIT>',<EOL>values=[observed_time_series, parameter_samples, num_steps_forecast]):<EOL><INDENT>[<EOL>observed_time_series,<EOL>mask<EOL>] = sts_util.canonicalize_observed_time_series_with_mask(<EOL>observed_time_series)<EOL>num_observed_steps = dist_util.prefer_static_value(<EOL>tf.shape(input=observed_time_series))[-<NUM_LIT:2>]<EOL>observed_data_ssm = model.make_state_space_model(<EOL>num_timesteps=num_observed_steps, param_vals=parameter_samples)<EOL>(_, _, _, predictive_means, predictive_covs, _, _<EOL>) = observed_data_ssm.forward_filter(observed_time_series, mask=mask)<EOL>parameter_samples = model._canonicalize_param_vals_as_map(parameter_samples)  <EOL>parameter_samples_with_reordered_batch_dimension = {<EOL>param.name: dist_util.move_dimension(<EOL>parameter_samples[param.name],<EOL><NUM_LIT:0>, -(<NUM_LIT:1> + _prefer_static_event_ndims(param.prior)))<EOL>for param in model.parameters}<EOL>forecast_prior = tfd.MultivariateNormalFullCovariance(<EOL>loc=dist_util.move_dimension(predictive_means[..., -<NUM_LIT:1>, :], <NUM_LIT:0>, -<NUM_LIT:2>),<EOL>covariance_matrix=dist_util.move_dimension(<EOL>predictive_covs[..., -<NUM_LIT:1>, :, :], <NUM_LIT:0>, -<NUM_LIT:3>))<EOL>kwargs = {}<EOL>if hasattr(model, '<STR_LIT>'):<EOL><INDENT>kwargs['<STR_LIT>'] = tf.convert_to_tensor(<EOL>value=model.constant_offset,<EOL>dtype=forecast_prior.dtype)[..., tf.newaxis]<EOL><DEDENT>forecast_ssm = model._make_state_space_model(  <EOL>num_timesteps=num_steps_forecast,<EOL>param_map=parameter_samples_with_reordered_batch_dimension,<EOL>initial_state_prior=forecast_prior,<EOL>initial_step=num_observed_steps,<EOL>**kwargs)<EOL>num_posterior_draws = dist_util.prefer_static_value(<EOL>forecast_ssm.batch_shape_tensor())[-<NUM_LIT:1>]<EOL>return tfd.MixtureSameFamily(<EOL>mixture_distribution=tfd.Categorical(<EOL>logits=tf.zeros([num_posterior_draws], dtype=forecast_ssm.dtype)),<EOL>components_distribution=forecast_ssm)<EOL><DEDENT>

Construct predictive distribution over future observations.

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

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

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

    #### Examples

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

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

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

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

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

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

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

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

      # OMITTED: take steps to optimize variational loss

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

    We can visualize the forecast by plotting:

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

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

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

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

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

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


top_row = tf.expand_dims(coefficients, -<NUM_LIT:2>)<EOL>coef_shape = dist_util.prefer_static_shape(coefficients)<EOL>batch_shape, order = coef_shape[:-<NUM_LIT:1>], coef_shape[-<NUM_LIT:1>]<EOL>remaining_rows = tf.concat([<EOL>tf.eye(order - <NUM_LIT:1>, dtype=coefficients.dtype, batch_shape=batch_shape),<EOL>tf.zeros(tf.concat([batch_shape, (order - <NUM_LIT:1>, <NUM_LIT:1>)], axis=<NUM_LIT:0>),<EOL>dtype=coefficients.dtype)<EOL>], axis=-<NUM_LIT:1>)<EOL>ar_matrix = tf.concat([top_row, remaining_rows], axis=-<NUM_LIT:2>)<EOL>return ar_matrix<EOL>

Build transition matrix for an autoregressive StateSpaceModel.

    When applied to a vector of previous values, this matrix computes
    the expected new value (summing the previous states according to the
    autoregressive coefficients) in the top dimension of the state space,
    and moves all previous values down by one dimension, 'forgetting' the
    final (least recent) value. That is, it looks like this:

    ```
    ar_matrix = [ coefs[0], coefs[1], ..., coefs[order]
                  1.,       0 ,       ..., 0.
                  0.,       1.,       ..., 0.
                  ...
                  0.,       0.,  ..., 1.,  0.            ]
    ```

    Args:
      coefficients: float `Tensor` of shape `concat([batch_shape, [order]])`.

    Returns:
      ar_matrix: float `Tensor` with shape `concat([batch_shape,
      [order, order]])`.


with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>',<EOL>values=[coefficients, level_scale, observation_noise_scale]) as name:<EOL><INDENT>dtype = initial_state_prior.dtype<EOL>coefficients = tf.convert_to_tensor(<EOL>value=coefficients, name='<STR_LIT>', dtype=dtype)<EOL>level_scale = tf.convert_to_tensor(<EOL>value=level_scale, name='<STR_LIT>', dtype=dtype)<EOL>observation_noise_scale = tf.convert_to_tensor(<EOL>value=observation_noise_scale,<EOL>name='<STR_LIT>', dtype=dtype)<EOL>order = tf.compat.dimension_value(coefficients.shape[-<NUM_LIT:1>])<EOL>if order is None:<EOL><INDENT>raise ValueError('<STR_LIT>')<EOL><DEDENT>self._order = order<EOL>self._coefficients = coefficients<EOL>self._level_scale = level_scale<EOL>super(AutoregressiveStateSpaceModel, self).__init__(<EOL>num_timesteps=num_timesteps,<EOL>transition_matrix=make_ar_transition_matrix(coefficients),<EOL>transition_noise=tfd.MultivariateNormalDiag(<EOL>scale_diag=tf.stack([level_scale] +<EOL>[tf.zeros_like(level_scale)] * (<EOL>self.order - <NUM_LIT:1>), axis=-<NUM_LIT:1>)),<EOL>observation_matrix=tf.concat([tf.ones([<NUM_LIT:1>, <NUM_LIT:1>], dtype=dtype),<EOL>tf.zeros([<NUM_LIT:1>, self.order - <NUM_LIT:1>],<EOL>dtype=dtype)],<EOL>axis=-<NUM_LIT:1>),<EOL>observation_noise=tfd.MultivariateNormalDiag(<EOL>scale_diag=observation_noise_scale[..., tf.newaxis]),<EOL>initial_state_prior=initial_state_prior,<EOL>initial_step=initial_step,<EOL>validate_args=validate_args,<EOL>name=name)<EOL><DEDENT>

Build a state space model implementing an autoregressive process.

        Args:
          num_timesteps: Scalar `int` `Tensor` number of timesteps to model
            with this distribution.
          coefficients: `float` `Tensor` of shape `concat(batch_shape, [order])`
            defining  the autoregressive coefficients. The coefficients are defined
            backwards in time: `coefficients[0] * level[t] + coefficients[1] *
            level[t-1] + ... + coefficients[order-1] * level[t-order+1]`.
          level_scale: Scalar (any additional dimensions are treated as batch
            dimensions) `float` `Tensor` indicating the standard deviation of the
            transition noise at each step.
          initial_state_prior: instance of `tfd.MultivariateNormal`
            representing the prior distribution on latent states.  Must have
            event shape `[order]`.
          observation_noise_scale: Scalar (any additional dimensions are
            treated as batch dimensions) `float` `Tensor` indicating the standard
            deviation of the observation noise.
            Default value: 0.
          initial_step: Optional scalar `int` `Tensor` specifying the starting
            timestep.
            Default value: 0.
          validate_args: Python `bool`. Whether to validate input
            with asserts. If `validate_args` is `False`, and the inputs are
            invalid, correct behavior is not guaranteed.
            Default value: `False`.
          name: Python `str` name prefixed to ops created by this class.
            Default value: "AutoregressiveStateSpaceModel".


with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>', values=[observed_time_series]) as name:<EOL><INDENT>masked_time_series = None<EOL>if observed_time_series is not None:<EOL><INDENT>masked_time_series = (<EOL>sts_util.canonicalize_observed_time_series_with_mask(<EOL>observed_time_series))<EOL><DEDENT>dtype = dtype_util.common_dtype(<EOL>[(masked_time_series.time_series<EOL>if masked_time_series is not None else None),<EOL>coefficients_prior,<EOL>level_scale_prior,<EOL>initial_state_prior], preferred_dtype=tf.float32)<EOL>if observed_time_series is not None:<EOL><INDENT>_, observed_stddev, observed_initial = sts_util.empirical_statistics(<EOL>masked_time_series)<EOL><DEDENT>else:<EOL><INDENT>observed_stddev, observed_initial = (<EOL>tf.convert_to_tensor(value=<NUM_LIT:1.>, dtype=dtype),<EOL>tf.convert_to_tensor(value=<NUM_LIT:0.>, dtype=dtype))<EOL><DEDENT>batch_ones = tf.ones(tf.concat([<EOL>tf.shape(input=observed_initial),  <EOL>[order]], axis=<NUM_LIT:0>), dtype=dtype)<EOL>if coefficients_prior is None:<EOL><INDENT>coefficients_prior = tfd.MultivariateNormalDiag(<EOL>scale_diag=batch_ones)<EOL><DEDENT>if level_scale_prior is None:<EOL><INDENT>level_scale_prior = tfd.LogNormal(<EOL>loc=tf.math.log(<NUM_LIT> *  observed_stddev), scale=<NUM_LIT>)<EOL><DEDENT>if (coefficients_prior.event_shape.is_fully_defined() and<EOL>order != coefficients_prior.event_shape[<NUM_LIT:0>]):<EOL><INDENT>raise ValueError("<STR_LIT>".format(<EOL>coefficients_prior.event_shape[<NUM_LIT:0>], order))<EOL><DEDENT>if initial_state_prior is None:<EOL><INDENT>initial_state_prior = tfd.MultivariateNormalDiag(<EOL>loc=observed_initial[..., tf.newaxis] * batch_ones,<EOL>scale_diag=(tf.abs(observed_initial) +<EOL>observed_stddev)[..., tf.newaxis] * batch_ones)<EOL><DEDENT>self._order = order<EOL>self._coefficients_prior = coefficients_prior<EOL>self._level_scale_prior = level_scale_prior<EOL>self._initial_state_prior = initial_state_prior<EOL>if coefficient_constraining_bijector is None:<EOL><INDENT>coefficient_constraining_bijector = tfb.Tanh()<EOL><DEDENT>super(Autoregressive, self).__init__(<EOL>parameters=[<EOL>Parameter('<STR_LIT>',<EOL>coefficients_prior,<EOL>coefficient_constraining_bijector),<EOL>Parameter('<STR_LIT>', level_scale_prior, tfb.Softplus())<EOL>],<EOL>latent_size=order,<EOL>name=name)<EOL><DEDENT>

Specify an autoregressive model.

        Args:
          order: scalar Python positive `int` specifying the number of past
            timesteps to regress on.
          coefficients_prior: optional `tfd.Distribution` instance specifying a
            prior on the `coefficients` parameter. If `None`, a default standard
            normal (`tfd.MultivariateNormalDiag(scale_diag=tf.ones([order]))`) prior
            is used.
            Default value: `None`.
          level_scale_prior: optional `tfd.Distribution` instance specifying a prior
            on the `level_scale` parameter. If `None`, a heuristic default prior is
            constructed based on the provided `observed_time_series`.
            Default value: `None`.
          initial_state_prior: optional `tfd.Distribution` instance specifying a
            prior on the initial state, corresponding to the values of the process
            at a set of size `order` of imagined timesteps before the initial step.
            If `None`, a heuristic default prior is constructed based on the
            provided `observed_time_series`.
            Default value: `None`.
          coefficient_constraining_bijector: optional `tfb.Bijector` instance
            representing a constraining mapping for the autoregressive coefficients.
            For example, `tfb.Tanh()` constrains the coefficients to lie in
            `(-1, 1)`, while `tfb.Softplus()` constrains them to be positive, and
            `tfb.Identity()` implies no constraint. If `None`, the default behavior
            constrains the coefficients to lie in `(-1, 1)` using a `Tanh` bijector.
            Default value: `None`.
          observed_time_series: optional `float` `Tensor` of shape
            `batch_shape + [T, 1]` (omitting the trailing unit dimension is also
            supported when `T > 1`), specifying an observed time series.
            Any priors not explicitly set will be given default values according to
            the scale of the observed time series (or batch of time series). May
            optionally be an instance of `tfp.sts.MaskedTimeSeries`, which includes
            a mask `Tensor` to specify timesteps with missing observations.
            Default value: `None`.
          name: the name of this model component.
            Default value: 'Autoregressive'.


def ar_explicit_logp(y, coefs, level_scale):

num_coefs = len(coefs)<EOL>lp = <NUM_LIT:0.><EOL>for i in range(num_coefs):<EOL><INDENT>zero_padded_y = np.zeros([num_coefs])<EOL>zero_padded_y[num_coefs - i:num_coefs] = y[:i]<EOL>pred_y = np.dot(zero_padded_y, coefs[::-<NUM_LIT:1>])<EOL>lp += tfd.Normal(pred_y, level_scale).log_prob(y[i])<EOL><DEDENT>for i in range(num_coefs, len(y)):<EOL><INDENT>pred_y = np.dot(y[i - num_coefs:i], coefs[::-<NUM_LIT:1>])<EOL>lp += tfd.Normal(pred_y, level_scale).log_prob(y[i])<EOL><DEDENT>return lp<EOL>

Manual log-prob computation for an autoregressive process.


ndarray = np.asarray(ndarray).astype(self.dtype)<EOL>return tf.compat.v1.placeholder_with_default(<EOL>input=ndarray, shape=ndarray.shape if self.use_static_shape else None)<EOL>

Convert a numpy array to a TF placeholder.

        Args:
          ndarray: any object convertible to a numpy array via `np.asarray()`.

        Returns:
          placeholder: a TensorFlow `placeholder` with default value given by the
          provided `ndarray`, dtype given by `self.dtype`, and shape specified
          statically only if `self.use_static_shape` is `True`.


<EOL>batch_shape = distributions[<NUM_LIT:0>].batch_shape<EOL>for distribution in distributions:<EOL><INDENT>batch_shape = tf.broadcast_static_shape(batch_shape,<EOL>distribution.batch_shape)<EOL><DEDENT>if batch_shape.is_fully_defined():<EOL><INDENT>return batch_shape.as_list()<EOL><DEDENT>batch_shape = distributions[<NUM_LIT:0>].batch_shape_tensor()<EOL>for distribution in distributions:<EOL><INDENT>batch_shape = tf.broadcast_dynamic_shape(batch_shape,<EOL>distribution.batch_shape_tensor())<EOL><DEDENT>return tf.convert_to_tensor(value=batch_shape)<EOL>

Get broadcast batch shape from distributions, statically if possible.


def pad_batch_dimension_for_multiple_chains(<EOL>observed_time_series, model, chain_batch_shape):

<EOL>[  <EOL>observed_time_series,<EOL>is_missing<EOL>] = canonicalize_observed_time_series_with_mask(observed_time_series)<EOL>event_ndims = <NUM_LIT:2>  <EOL>model_batch_ndims = (<EOL>model.batch_shape.ndims if model.batch_shape.ndims is not None else<EOL>tf.shape(input=model.batch_shape_tensor())[<NUM_LIT:0>])<EOL>chain_batch_shape = tf.convert_to_tensor(<EOL>value=chain_batch_shape, name='<STR_LIT>', dtype=tf.int32)<EOL>if not chain_batch_shape.shape.is_fully_defined():<EOL><INDENT>raise ValueError('<STR_LIT>'.format(<EOL>chain_batch_shape))<EOL><DEDENT>if chain_batch_shape.shape.ndims == <NUM_LIT:0>:  <EOL><INDENT>chain_batch_shape = chain_batch_shape[tf.newaxis]<EOL><DEDENT>chain_batch_ndims = tf.compat.dimension_value(chain_batch_shape.shape[<NUM_LIT:0>])<EOL>def do_padding(observed_time_series_tensor):<EOL><INDENT>current_sample_shape = tf.shape(<EOL>input=observed_time_series_tensor)[:-(model_batch_ndims + event_ndims)]<EOL>current_batch_and_event_shape = tf.shape(<EOL>input=observed_time_series_tensor)[-(model_batch_ndims + event_ndims):]<EOL>return tf.reshape(<EOL>tensor=observed_time_series_tensor,<EOL>shape=tf.concat([<EOL>current_sample_shape,<EOL>tf.ones([chain_batch_ndims], dtype=tf.int32),<EOL>current_batch_and_event_shape], axis=<NUM_LIT:0>))<EOL><DEDENT>observed_time_series = prefer_static.cond(<EOL>(dist_util.prefer_static_rank(observed_time_series) ><EOL>model_batch_ndims + event_ndims),<EOL>lambda: do_padding(observed_time_series),<EOL>lambda: observed_time_series)<EOL>if is_missing is not None:<EOL><INDENT>is_missing = prefer_static.cond(<EOL>(dist_util.prefer_static_rank(is_missing) ><EOL>model_batch_ndims + event_ndims),<EOL>lambda: do_padding(is_missing),<EOL>lambda: is_missing)<EOL>return missing_values_util.MaskedTimeSeries(observed_time_series,<EOL>is_missing=is_missing)<EOL><DEDENT>return observed_time_series<EOL>

Expand the observed time series with extra batch dimension(s).


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

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<EOL>for tensor in distribution._graph_parents]  <EOL>with tf.compat.v1.name_scope('<STR_LIT>', values=graph_parents):<EOL><INDENT>if all([isinstance(mvn, tfd.MultivariateNormalDiag)<EOL>for mvn in distributions]):<EOL><INDENT>return tfd.MultivariateNormalDiag(<EOL>loc=sum([mvn.mean() for mvn in distributions]),<EOL>scale_diag=tf.sqrt(sum([<EOL>mvn.scale.diag**<NUM_LIT:2> for mvn in distributions])))<EOL><DEDENT>else:<EOL><INDENT>raise NotImplementedError(<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'.format(distributions))<EOL><DEDENT><DEDENT>

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.


with tf.compat.v1.name_scope(<EOL>'<STR_LIT>', values=[observed_time_series]):<EOL><INDENT>[<EOL>observed_time_series,<EOL>mask<EOL>] = canonicalize_observed_time_series_with_mask(observed_time_series)<EOL>squeezed_series = observed_time_series[..., <NUM_LIT:0>]<EOL>if mask is None:<EOL><INDENT>observed_mean, observed_variance = tf.nn.moments(<EOL>x=squeezed_series, axes=-<NUM_LIT:1>)<EOL>observed_initial = squeezed_series[..., <NUM_LIT:0>]<EOL><DEDENT>else:<EOL><INDENT>broadcast_mask = tf.broadcast_to(tf.cast(mask, tf.bool),<EOL>tf.shape(input=squeezed_series))<EOL>observed_mean, observed_variance = (<EOL>missing_values_util.moments_of_masked_time_series(<EOL>squeezed_series, broadcast_mask=broadcast_mask))<EOL>try:<EOL><INDENT>observed_initial = (<EOL>missing_values_util.initial_value_of_masked_time_series(<EOL>squeezed_series, broadcast_mask=broadcast_mask))<EOL><DEDENT>except NotImplementedError:<EOL><INDENT>tf.compat.v1.logging.warn(<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'<EOL>'<STR_LIT>')<EOL>observed_initial = observed_mean<EOL><DEDENT><DEDENT>observed_stddev = tf.sqrt(observed_variance)<EOL>observed_initial_centered = observed_initial - observed_mean<EOL>return observed_mean, observed_stddev, observed_initial_centered<EOL><DEDENT>

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(<EOL>'<STR_LIT>', values=[observed_time_series_tensor]):<EOL><INDENT>if (observed_time_series_tensor.shape.ndims is not None and<EOL>tf.compat.dimension_value(<EOL>observed_time_series_tensor.shape[-<NUM_LIT:1>]) is not None):<EOL><INDENT>expanded_time_series = (<EOL>observed_time_series_tensor<EOL>if observed_time_series_tensor.shape[-<NUM_LIT:1>] == <NUM_LIT:1><EOL>else observed_time_series_tensor[..., tf.newaxis])<EOL><DEDENT>else:<EOL><INDENT>expanded_time_series = tf.cond(<EOL>pred=tf.equal(tf.shape(input=observed_time_series_tensor)[-<NUM_LIT:1>], <NUM_LIT:1>),<EOL>true_fn=lambda: observed_time_series_tensor,<EOL>false_fn=lambda: observed_time_series_tensor[..., tf.newaxis])<EOL><DEDENT>return expanded_time_series<EOL><DEDENT>

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('<STR_LIT>'):<EOL><INDENT>if hasattr(maybe_masked_observed_time_series, '<STR_LIT>'):<EOL><INDENT>observed_time_series = (<EOL>maybe_masked_observed_time_series.time_series)<EOL>is_missing = maybe_masked_observed_time_series.is_missing<EOL><DEDENT>else:<EOL><INDENT>observed_time_series = maybe_masked_observed_time_series<EOL>is_missing = None<EOL><DEDENT>observed_time_series = tf.convert_to_tensor(value=observed_time_series,<EOL>name='<STR_LIT>')<EOL>observed_time_series = _maybe_expand_trailing_dim(observed_time_series)<EOL>if is_missing is not None:<EOL><INDENT>is_missing = tf.convert_to_tensor(<EOL>value=is_missing, name='<STR_LIT>', dtype_hint=tf.bool)<EOL><DEDENT>return missing_values_util.MaskedTimeSeries(observed_time_series,<EOL>is_missing=is_missing)<EOL><DEDENT>

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


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

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]`.


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

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_timesteps = tf.shape(input=time_series_tensor)[-<NUM_LIT:1>]<EOL>unmasked_negindices = (<EOL>tf.cast(~broadcast_mask, tf.int32) *<EOL>tf.range(num_timesteps, <NUM_LIT:0>, -<NUM_LIT:1>))<EOL>first_unmasked_indices = num_timesteps - tf.reduce_max(<EOL>input_tensor=unmasked_negindices, axis=-<NUM_LIT:1>)<EOL>if first_unmasked_indices.shape.ndims is None:<EOL><INDENT>raise NotImplementedError(<EOL>'<STR_LIT>'<EOL>'<STR_LIT>')  <EOL><DEDENT>return tf.squeeze(tf.compat.v1.batch_gather(<EOL>params=time_series_tensor,<EOL>indices=first_unmasked_indices[..., tf.newaxis]), axis=-<NUM_LIT:1>)<EOL>

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


def _build_tensor(self, ndarray, dtype=None):

ndarray = np.asarray(ndarray).astype(self.dtype if dtype is None else dtype)<EOL>return tf.compat.v1.placeholder_with_default(<EOL>input=ndarray, shape=ndarray.shape if self.use_static_shape else None)<EOL>

Convert a numpy array to a TF placeholder.

        Args:
          ndarray: any object convertible to a numpy array via `np.asarray()`.
          dtype: optional `dtype`.

        Returns:
          placeholder: a TensorFlow `placeholder` with default value given by the
          provided `ndarray`, dtype given by `self.dtype` (if not specified), and
          shape specified statically only if `self.use_static_shape` is `True`.


def _build_placeholder(self, ndarray, dtype=None):

if dtype is None:<EOL><INDENT>dtype = self.dtype<EOL><DEDENT>ndarray = np.asarray(ndarray).astype(dtype)<EOL>return tf.compat.v1.placeholder_with_default(<EOL>input=ndarray, shape=ndarray.shape if self.use_static_shape else None)<EOL>

Convert a numpy array to a TF placeholder.

        Args:
          ndarray: any object convertible to a numpy array via `np.asarray()`.
          dtype: optional `dtype`; if not specified, defaults to `self.dtype`.

        Returns:
          placeholder: a TensorFlow `placeholder` with default value given by the
          provided `ndarray`, dtype given by `self.dtype`, and shape specified
          statically only if `self.use_static_shape` is `True`.


dummy_mvndiag = tfd.MultivariateNormalDiag(<EOL>scale_diag=tf.ones([<NUM_LIT:0>], dtype=dtype))<EOL>dummy_mvndiag.covariance = lambda: dummy_mvndiag.variance()[..., tf.newaxis]<EOL>return dummy_mvndiag<EOL>

Build a zero-dimensional MVNDiag object.


def observation_noise_fn(t):<EOL><INDENT>current_slice = timeseries[..., t, :]<EOL>return tfd.MultivariateNormalDiag(<EOL>loc=current_slice,<EOL>scale_diag=tf.zeros_like(current_slice))<EOL><DEDENT>return observation_noise_fn<EOL>

Build an observation_noise_fn that observes a Tensor timeseries.


with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>', values=[design_matrix]) as name:<EOL><INDENT>if not isinstance(design_matrix, tfl.LinearOperator):<EOL><INDENT>design_matrix = tfl.LinearOperatorFullMatrix(<EOL>tf.convert_to_tensor(value=design_matrix, name='<STR_LIT>'),<EOL>name='<STR_LIT>')<EOL><DEDENT>if tf.compat.dimension_value(design_matrix.shape[-<NUM_LIT:1>]) is not None:<EOL><INDENT>num_features = design_matrix.shape[-<NUM_LIT:1>]<EOL><DEDENT>else:<EOL><INDENT>num_features = design_matrix.shape_tensor()[-<NUM_LIT:1>]<EOL><DEDENT>if weights_prior is None:<EOL><INDENT>weights_prior = tfd.StudentT(<EOL>df=<NUM_LIT:5>,<EOL>loc=tf.zeros([], dtype=design_matrix.dtype),<EOL>scale=<NUM_LIT:10> * tf.ones([], dtype=design_matrix.dtype))<EOL><DEDENT>if weights_prior.event_shape.ndims == <NUM_LIT:0>:<EOL><INDENT>if design_matrix.batch_shape.is_fully_defined():<EOL><INDENT>design_matrix_batch_shape_ = design_matrix.batch_shape<EOL><DEDENT>else:<EOL><INDENT>design_matrix_batch_shape_ = design_matrix.batch_shape_tensor()<EOL><DEDENT>weights_prior = tfd.TransformedDistribution(<EOL>weights_prior,<EOL>bijector=tfb.Identity(),<EOL>batch_shape=design_matrix_batch_shape_,<EOL>event_shape=[num_features])<EOL><DEDENT>tf.debugging.assert_same_float_dtype([design_matrix, weights_prior])<EOL>self._design_matrix = design_matrix<EOL>super(LinearRegression, self).__init__(<EOL>parameters=[<EOL>Parameter('<STR_LIT>', weights_prior, tfb.Identity()),<EOL>],<EOL>latent_size=<NUM_LIT:0>,<EOL>name=name)<EOL><DEDENT>

Specify a linear regression model.

        Note: the statistical behavior of the regression is determined by
        the broadcasting behavior of the `weights` `Tensor`:

        * `weights_prior.batch_shape == []`: shares a single set of weights across
          all design matrices and observed time series. This may make sense if
          the features in each design matrix have the same semantics (e.g.,
          grouping observations by country, with per-country design matrices
          capturing the same set of national economic indicators per country).
        * `weights_prior.batch_shape == `design_matrix.batch_shape`: fits separate
          weights for each design matrix. If there are multiple observed time series
          for each design matrix, this shares statistical strength over those
          observations.
        * `weights_prior.batch_shape == `observed_time_series.batch_shape`: fits a
          separate regression for each individual time series.

        When modeling batches of time series, you should think carefully about
        which behavior makes sense, and specify `weights_prior` accordingly:
        the defaults may not do what you want!

        Args:
          design_matrix: float `Tensor` of shape `concat([batch_shape,
            [num_timesteps, num_features]])`. This may also optionally be
            an instance of `tf.linalg.LinearOperator`.
          weights_prior: `tfd.Distribution` representing a prior over the regression
            weights. Must have event shape `[num_features]` and batch shape
            broadcastable to the design matrix's `batch_shape`. Alternately,
            `event_shape` may be scalar (`[]`), in which case the prior is
            internally broadcast as `TransformedDistribution(weights_prior,
            tfb.Identity(), event_shape=[num_features],
            batch_shape=design_matrix.batch_shape)`. If `None`,
            defaults to `StudentT(df=5, loc=0., scale=10.)`, a weakly-informative
            prior loosely inspired by the [Stan prior choice recommendations](
            https://github.com/stan-dev/stan/wiki/Prior-Choice-Recommendations).
            Default value: `None`.
          name: the name of this model component.
            Default value: 'LinearRegression'.


return self._design_matrix<EOL>

LinearOperator representing the design matrix.


with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>',<EOL>values=[design_matrix, weights_prior_scale]) as name:<EOL><INDENT>if not isinstance(design_matrix, tfl.LinearOperator):<EOL><INDENT>design_matrix = tfl.LinearOperatorFullMatrix(<EOL>tf.convert_to_tensor(value=design_matrix, name='<STR_LIT>'),<EOL>name='<STR_LIT>')<EOL><DEDENT>if tf.compat.dimension_value(design_matrix.shape[-<NUM_LIT:1>]) is not None:<EOL><INDENT>num_features = design_matrix.shape[-<NUM_LIT:1>]<EOL><DEDENT>else:<EOL><INDENT>num_features = design_matrix.shape_tensor()[-<NUM_LIT:1>]<EOL><DEDENT>if weights_batch_shape is None:<EOL><INDENT>weights_batch_shape = design_matrix.batch_shape_tensor()<EOL><DEDENT>else:<EOL><INDENT>weights_batch_shape = tf.convert_to_tensor(value=weights_batch_shape,<EOL>dtype=tf.int32)<EOL><DEDENT>weights_shape = tf.concat([weights_batch_shape, [num_features]], axis=<NUM_LIT:0>)<EOL>dtype = design_matrix.dtype<EOL>self._design_matrix = design_matrix<EOL>self._weights_prior_scale = weights_prior_scale<EOL>ones_like_weights_batch = tf.ones(weights_batch_shape, dtype=dtype)<EOL>ones_like_weights = tf.ones(weights_shape, dtype=dtype)<EOL>super(SparseLinearRegression, self).__init__(<EOL>parameters=[<EOL>Parameter('<STR_LIT>',<EOL>prior=tfd.InverseGamma(<EOL><NUM_LIT:0.5> * ones_like_weights_batch,<EOL><NUM_LIT:0.5> * ones_like_weights_batch),<EOL>bijector=tfb.Softplus()),<EOL>Parameter('<STR_LIT>',<EOL>prior=tfd.HalfNormal(<EOL>scale=ones_like_weights_batch),<EOL>bijector=tfb.Softplus()),<EOL>Parameter('<STR_LIT>',<EOL>prior=tfd.Independent(tfd.InverseGamma(<EOL><NUM_LIT:0.5> * ones_like_weights,<EOL><NUM_LIT:0.5> * ones_like_weights),<EOL>reinterpreted_batch_ndims=<NUM_LIT:1>),<EOL>bijector=tfb.Softplus()),<EOL>Parameter('<STR_LIT>',<EOL>prior=tfd.Independent(tfd.HalfNormal(<EOL>scale=ones_like_weights),<EOL>reinterpreted_batch_ndims=<NUM_LIT:1>),<EOL>bijector=tfb.Softplus()),<EOL>Parameter('<STR_LIT>',<EOL>prior=tfd.Independent(tfd.Normal(<EOL>loc=tf.zeros_like(ones_like_weights),<EOL>scale=ones_like_weights),<EOL>reinterpreted_batch_ndims=<NUM_LIT:1>),<EOL>bijector=tfb.Identity())<EOL>],<EOL>latent_size=<NUM_LIT:0>,<EOL>name=name)<EOL><DEDENT>

Specify a sparse linear regression model.

        Args:
          design_matrix: float `Tensor` of shape `concat([batch_shape,
            [num_timesteps, num_features]])`. This may also optionally be
            an instance of `tf.linalg.LinearOperator`.
          weights_prior_scale: float `Tensor` defining the scale of the Horseshoe
            prior on regression weights. Small values encourage the weights to be
            sparse. The shape must broadcast with `weights_batch_shape`.
            Default value: `0.1`.
          weights_batch_shape: if `None`, defaults to
            `design_matrix.batch_shape_tensor()`. Must broadcast with the batch
            shape of `design_matrix`.
            Default value: `None`.
          name: the name of this model component.
            Default value: 'SparseLinearRegression'.


return self._design_matrix<EOL>


global_scale = (global_scale_noncentered *<EOL>tf.sqrt(global_scale_variance) *<EOL>self.weights_prior_scale)<EOL>local_scales = local_scales_noncentered * tf.sqrt(local_scale_variances)<EOL>return weights_noncentered * local_scales * global_scale[..., tf.newaxis]<EOL>

Build regression weights from model parameters.


ndarray = np.asarray(ndarray).astype(self.dtype)<EOL>return tf.compat.v1.placeholder_with_default(<EOL>input=ndarray, shape=ndarray.shape if self.use_static_shape else None)<EOL>

Convert a numpy array to a TF placeholder.

        Args:
          ndarray: any object convertible to a numpy array via `np.asarray()`.

        Returns:
          placeholder: a TensorFlow `placeholder` with default value given by the
          provided `ndarray`, dtype given by `self.dtype`, and shape specified
          statically only if `self.use_static_shape` is `True`.


unconstrained_prior_sample = parameter.bijector.inverse(<EOL>parameter.prior.sample(init_sample_shape, seed=seed))<EOL>uniform_initializer = <NUM_LIT:4> * tf.random.uniform(<EOL>tf.shape(input=unconstrained_prior_sample),<EOL>dtype=unconstrained_prior_sample.dtype,<EOL>seed=seed) - <NUM_LIT:2><EOL>if return_constrained:<EOL><INDENT>return parameter.bijector.forward(uniform_initializer)<EOL><DEDENT>else:<EOL><INDENT>return uniform_initializer<EOL><DEDENT>

Initialize from a uniform [-2, 2] distribution in unconstrained space.

    Args:
      parameter: `sts.Parameter` named tuple instance.
      return_constrained: if `True`, re-applies the constraining bijector
        to return initializations in the original domain. Otherwise, returns
        initializations in the unconstrained space.
        Default value: `True`.
      init_sample_shape: `sample_shape` of the sampled initializations.
        Default value: `[]`.
      seed: Python integer to seed the random number generator.

    Returns:
      uniform_initializer: `Tensor` of shape `concat([init_sample_shape,
      parameter.prior.batch_shape, transformed_event_shape])`, where
      `transformed_event_shape` is `parameter.prior.event_shape`, if
      `return_constrained=True`, and otherwise it is
      `parameter.bijector.inverse_event_shape(parameteter.prior.event_shape)`.


loc = tf.compat.v1.get_variable(<EOL>param.name + '<STR_LIT>',<EOL>initializer=lambda: initial_loc_fn(param),<EOL>dtype=param.prior.dtype,<EOL>use_resource=True)<EOL>scale = tf.nn.softplus(<EOL>tf.compat.v1.get_variable(<EOL>param.name + '<STR_LIT>',<EOL>initializer=lambda: -<NUM_LIT:4> * tf.ones_like(initial_loc_fn(param)),<EOL>dtype=param.prior.dtype,<EOL>use_resource=True))<EOL>q = tfd.Normal(loc=loc, scale=scale)<EOL>if (param.prior.event_shape.ndims is None<EOL>or param.prior.event_shape.ndims > <NUM_LIT:0>):<EOL><INDENT>q = tfd.Independent(<EOL>q, reinterpreted_batch_ndims=param.prior.event_shape.ndims)<EOL><DEDENT>return tfd.TransformedDistribution(q, param.bijector)<EOL>

Built a transformed-normal variational dist over a parameter's support.


with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>',<EOL>values=[observed_time_series]) as name:<EOL><INDENT>seed = tfd.SeedStream(<EOL>seed, salt='<STR_LIT>')<EOL>variational_distributions = collections.OrderedDict()<EOL>variational_samples = []<EOL>for param in model.parameters:<EOL><INDENT>def initial_loc_fn(param):<EOL><INDENT>return sample_uniform_initial_state(<EOL>param, return_constrained=True,<EOL>init_sample_shape=init_batch_shape,<EOL>seed=seed())<EOL><DEDENT>q = _build_trainable_posterior(param, initial_loc_fn=initial_loc_fn)<EOL>variational_distributions[param.name] = q<EOL>variational_samples.append(q.sample(seed=seed()))<EOL><DEDENT>observed_time_series = sts_util.pad_batch_dimension_for_multiple_chains(<EOL>observed_time_series, model, chain_batch_shape=init_batch_shape)<EOL>log_prob_fn = model.joint_log_prob(observed_time_series)<EOL>expected_log_joint = log_prob_fn(*variational_samples)<EOL>entropy = tf.reduce_sum(<EOL>input_tensor=[<EOL>-q.log_prob(sample) for (q, sample) in zip(<EOL>variational_distributions.values(), variational_samples)<EOL>],<EOL>axis=<NUM_LIT:0>)<EOL>variational_loss = -(expected_log_joint + entropy)  <EOL><DEDENT>return variational_loss, variational_distributions<EOL>

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


def _minimize_in_graph(build_loss_fn, num_steps=<NUM_LIT:200>, optimizer=None):

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

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


with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>', values=[observed_time_series]) as name:<EOL><INDENT>seed = tfd.SeedStream(seed, salt='<STR_LIT>')<EOL>if initial_step_size is None or initial_state is None:<EOL><INDENT>def make_variational():<EOL><INDENT>return build_factored_variational_loss(<EOL>model, observed_time_series,<EOL>init_batch_shape=chain_batch_shape, seed=seed())<EOL><DEDENT>make_variational = tf.compat.v1.make_template('<STR_LIT>',<EOL>make_variational)<EOL>_, variational_distributions = make_variational()<EOL>minimize_op = _minimize_in_graph(<EOL>build_loss_fn=lambda: make_variational()[<NUM_LIT:0>],  <EOL>num_steps=num_variational_steps,<EOL>optimizer=variational_optimizer)<EOL>with tf.control_dependencies([minimize_op]):<EOL><INDENT>if initial_state is None:<EOL><INDENT>initial_state = [tf.stop_gradient(d.sample())<EOL>for d in variational_distributions.values()]<EOL><DEDENT>if initial_step_size is None:<EOL><INDENT>initial_step_size = [<EOL>transformed_q.distribution.stddev()<EOL>for transformed_q in variational_distributions.values()]<EOL><DEDENT><DEDENT><DEDENT>observed_time_series = sts_util.pad_batch_dimension_for_multiple_chains(<EOL>observed_time_series, model, chain_batch_shape=chain_batch_shape)<EOL>samples, kernel_results = mcmc.sample_chain(<EOL>num_results=num_results,<EOL>current_state=initial_state,<EOL>num_burnin_steps=num_warmup_steps,<EOL>kernel=mcmc.SimpleStepSizeAdaptation(<EOL>inner_kernel=mcmc.TransformedTransitionKernel(<EOL>inner_kernel=mcmc.HamiltonianMonteCarlo(<EOL>target_log_prob_fn=model.joint_log_prob(<EOL>observed_time_series),<EOL>step_size=initial_step_size,<EOL>num_leapfrog_steps=num_leapfrog_steps,<EOL>state_gradients_are_stopped=True,<EOL>seed=seed()),<EOL>bijector=[param.bijector for param in model.parameters]),<EOL>num_adaptation_steps=int(num_warmup_steps * <NUM_LIT>),<EOL>adaptation_rate=tf.convert_to_tensor(<EOL>value=<NUM_LIT:0.1>, dtype=initial_state[<NUM_LIT:0>].dtype)),<EOL>parallel_iterations=<NUM_LIT:1> if seed is not None else <NUM_LIT:10>)<EOL>return samples, kernel_results<EOL><DEDENT>

Draw posterior samples using Hamiltonian Monte Carlo (HMC).

    Markov chain Monte Carlo (MCMC) methods are considered the gold standard of
    Bayesian inference; under suitable conditions and in the limit of infinitely
    many draws they generate samples from the true posterior distribution. HMC [1]
    uses gradients of the model's log-density function to propose samples,
    allowing it to exploit posterior geometry. However, it is computationally more
    expensive than variational inference and relatively sensitive to tuning.

    This method attempts to provide a sensible default approach for fitting
    StructuralTimeSeries models using HMC. It first runs variational inference as
    a fast posterior approximation, and initializes the HMC sampler from the
    variational posterior, using the posterior standard deviations to set
    per-variable step sizes (equivalently, a diagonal mass matrix). During the
    warmup phase, it adapts the step size to target an acceptance rate of 0.75,
    which is thought to be in the desirable range for optimal mixing [2].

    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.
      num_results: Integer number of Markov chain draws.
        Default value: `100`.
      num_warmup_steps: Integer number of steps to take before starting to
        collect results. The warmup steps are also used to adapt the step size
        towards a target acceptance rate of 0.75.
        Default value: `50`.
      num_leapfrog_steps: Integer number of steps to run the leapfrog integrator
        for. Total progress per HMC step is roughly proportional to
        `step_size * num_leapfrog_steps`.
        Default value: `15`.
      initial_state: Optional Python `list` of `Tensor`s, one for each model
        parameter, representing the initial state(s) of the Markov chain(s). These
        should have shape `concat([chain_batch_shape, param.prior.batch_shape,
        param.prior.event_shape])`. If `None`, the initial state is set
        automatically using a sample from a variational posterior.
        Default value: `None`.
      initial_step_size: Python `list` of `Tensor`s, one for each model parameter,
        representing the step size for the leapfrog integrator. Must
        broadcast with the shape of `initial_state`. Larger step sizes lead to
        faster progress, but too-large step sizes make rejection exponentially
        more likely. If `None`, the step size is set automatically using the
        standard deviation of a variational posterior.
        Default value: `None`.
      chain_batch_shape: Batch shape (Python `tuple`, `list`, or `int`) of chains
        to run in parallel.
        Default value: `[]` (i.e., a single chain).
      num_variational_steps: Python `int` number of steps to run the variational
        optimization to determine the initial state and step sizes.
        Default value: `150`.
      variational_optimizer: Optional `tf.train.Optimizer` instance to use in
        the variational optimization. If `None`, defaults to
        `tf.train.AdamOptimizer(0.1)`.
        Default value: `None`.
      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., 'fit_with_hmc').

    Returns:
      samples: Python `list` of `Tensors` representing posterior samples of model
        parameters, with shapes `[concat([[num_results], chain_batch_shape,
        param.prior.batch_shape, param.prior.event_shape]) for param in
        model.parameters]`.
      kernel_results: A (possibly nested) `tuple`, `namedtuple` or `list` of
        `Tensor`s representing internal calculations made within the HMC sampler.

    #### 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 draw posterior samples using HMC under default settings:

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

    with tf.Session() as sess:
      sess.run(tf.global_variables_initializer())
      samples_, kernel_results_ = sess.run((samples, kernel_results))

    print("acceptance rate: {}".format(
      np.mean(kernel_results_.inner_results.is_accepted, axis=0)))
    print("posterior means: {}".format(
      {param.name: np.mean(param_draws, axis=0)
       for (param, param_draws) in zip(model.parameters, samples_)}))
    ```

    We can also run multiple chains. This may help diagnose convergence issues
    and allows us to exploit vectorization to draw samples more quickly, although
    warmup still requires the same number of sequential steps.

    ```python
    from matplotlib import pylab as plt

    samples, kernel_results = tfp.sts.fit_with_hmc(
      model, observed_time_series, chain_batch_shape=[10])

    with tf.Session() as sess:
      sess.run(tf.global_variables_initializer())
      samples_, kernel_results_ = sess.run((samples, kernel_results))

    print("acceptance rate: {}".format(
      np.mean(kernel_results_.inner_results.inner_results.is_accepted, axis=0)))

    # Plot the sampled traces for each parameter. If the chains have mixed, their
    # traces should all cover the same region of state space, frequently crossing
    # over each other.
    for (param, param_draws) in zip(model.parameters, samples_):
      if param.prior.event_shape.ndims > 0:
        print("Only plotting traces for scalar parameters, skipping {}".format(
          param.name))
        continue
      plt.figure(figsize=[10, 4])
      plt.title(param.name)
      plt.plot(param_draws)
      plt.ylabel(param.name)
      plt.xlabel("HMC step")

    # Combining the samples from multiple chains into a single dimension allows
    # us to easily pass sampled parameters to downstream forecasting methods.
    combined_samples_ = [np.reshape(param_draws,
                                    [-1] + list(param_draws.shape[2:]))
                         for param_draws in samples_]
    ```

    For greater flexibility, you may prefer to implement your own sampler using
    the TensorFlow Probability primitives in `tfp.mcmc`. The following recipe
    constructs a basic HMC sampler, using a `TransformedTransitionKernel` to
    incorporate constraints on the parameter space.

    ```python
    transformed_hmc_kernel = mcmc.TransformedTransitionKernel(
        inner_kernel=mcmc.SimpleStepSizeAdaptation(
            inner_kernel=mcmc.HamiltonianMonteCarlo(
                target_log_prob_fn=model.joint_log_prob(observed_time_series),
                step_size=step_size,
                num_leapfrog_steps=num_leapfrog_steps,
                state_gradients_are_stopped=True,
                seed=seed),
            num_adaptation_steps = int(0.8 * num_warmup_steps)),
        bijector=[param.bijector for param in model.parameters])

    # Initialize from a Uniform[-2, 2] distribution in unconstrained space.
    initial_state = [tfp.sts.sample_uniform_initial_state(
      param, return_constrained=True) for param in model.parameters]

    samples, kernel_results = tfp.mcmc.sample_chain(
      kernel=transformed_hmc_kernel,
      num_results=num_results,
      current_state=initial_state,
      num_burnin_steps=num_warmup_steps)
    ```

    #### References

    [1]: Radford Neal. MCMC Using Hamiltonian Dynamics. _Handbook of Markov Chain
         Monte Carlo_, 2011. https://arxiv.org/abs/1206.1901
    [2]  M.J. Betancourt, Simon Byrne, and Mark Girolami. Optimizing The
         Integrator Step Size for Hamiltonian Monte Carlo.
         https://arxiv.org/abs/1411.6669


with tf.compat.v1.name_scope(name, '<STR_LIT>',<EOL>[level_scale]) as name:<EOL><INDENT>dtype = initial_state_prior.dtype<EOL>level_scale = tf.convert_to_tensor(<EOL>value=level_scale, name='<STR_LIT>', dtype=dtype)<EOL>observation_noise_scale = tf.convert_to_tensor(<EOL>value=observation_noise_scale,<EOL>name='<STR_LIT>',<EOL>dtype=dtype)<EOL>self._level_scale = level_scale<EOL>self._observation_noise_scale = observation_noise_scale<EOL>super(LocalLevelStateSpaceModel, self).__init__(<EOL>num_timesteps=num_timesteps,<EOL>transition_matrix=tf.constant(<EOL>[[<NUM_LIT:1.>]], dtype=dtype, name='<STR_LIT>'),<EOL>transition_noise=tfd.MultivariateNormalDiag(<EOL>scale_diag=level_scale[..., tf.newaxis], name='<STR_LIT>'),<EOL>observation_matrix=tf.constant(<EOL>[[<NUM_LIT:1.>]], dtype=dtype, name='<STR_LIT>'),<EOL>observation_noise=tfd.MultivariateNormalDiag(<EOL>scale_diag=observation_noise_scale[..., tf.newaxis],<EOL>name='<STR_LIT>'),<EOL>initial_state_prior=initial_state_prior,<EOL>initial_step=initial_step,<EOL>allow_nan_stats=allow_nan_stats,<EOL>validate_args=validate_args,<EOL>name=name)<EOL><DEDENT>

Build a state space model implementing a local level.

        Args:
          num_timesteps: Scalar `int` `Tensor` number of timesteps to model
            with this distribution.
          level_scale: Scalar (any additional dimensions are treated as batch
            dimensions) `float` `Tensor` indicating the standard deviation of the
            level transitions.
          initial_state_prior: instance of `tfd.MultivariateNormal`
            representing the prior distribution on latent states.  Must have
            event shape `[1]` (as `tfd.LinearGaussianStateSpaceModel` requires a
            rank-1 event shape).
          observation_noise_scale: Scalar (any additional dimensions are
            treated as batch dimensions) `float` `Tensor` indicating the standard
            deviation of the observation noise.
          initial_step: Optional scalar `int` `Tensor` specifying the starting
            timestep.
            Default value: 0.
          validate_args: Python `bool`. Whether to validate input
            with asserts. If `validate_args` is `False`, and the inputs are
            invalid, correct behavior is not guaranteed.
            Default value: `False`.
          allow_nan_stats: Python `bool`. If `False`, raise an
            exception if a statistic (e.g. mean/mode/etc...) is undefined for any
            batch member. If `True`, batch members with valid parameters leading to
            undefined statistics will return NaN for this statistic.
            Default value: `True`.
          name: Python `str` name prefixed to ops created by this class.
            Default value: "LocalLevelStateSpaceModel".


return self._level_scale<EOL>

Standard deviation of the level transitions.


return self._observation_noise_scale<EOL>

Standard deviation of the observation noise.


with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>', values=[observed_time_series]) as name:<EOL><INDENT>dtype = dtype_util.common_dtype([level_scale_prior, initial_level_prior])<EOL>if level_scale_prior is None or initial_level_prior is None:<EOL><INDENT>if observed_time_series is not None:<EOL><INDENT>_, observed_stddev, observed_initial = (<EOL>sts_util.empirical_statistics(observed_time_series))<EOL><DEDENT>else:<EOL><INDENT>observed_stddev, observed_initial = (tf.convert_to_tensor(<EOL>value=<NUM_LIT:1.>, dtype=dtype), tf.convert_to_tensor(<EOL>value=<NUM_LIT:0.>, dtype=dtype))<EOL><DEDENT><DEDENT>if level_scale_prior is None:<EOL><INDENT>level_scale_prior = tfd.LogNormal(<EOL>loc=tf.math.log(<NUM_LIT> * observed_stddev),<EOL>scale=<NUM_LIT>,<EOL>name='<STR_LIT>')<EOL><DEDENT>if initial_level_prior is None:<EOL><INDENT>self._initial_state_prior = tfd.MultivariateNormalDiag(<EOL>loc=observed_initial[..., tf.newaxis],<EOL>scale_diag=(<EOL>tf.abs(observed_initial) + observed_stddev)[..., tf.newaxis],<EOL>name='<STR_LIT>')<EOL><DEDENT>else:<EOL><INDENT>self._initial_state_prior = tfd.MultivariateNormalDiag(<EOL>loc=initial_level_prior.mean()[..., tf.newaxis],<EOL>scale_diag=initial_level_prior.stddev()[..., tf.newaxis])<EOL><DEDENT>super(LocalLevel, self).__init__(<EOL>parameters=[<EOL>Parameter('<STR_LIT>', level_scale_prior, tfb.Softplus()),<EOL>],<EOL>latent_size=<NUM_LIT:1>,<EOL>name=name)<EOL><DEDENT>

Specify a local level model.

        Args:
          level_scale_prior: optional `tfd.Distribution` instance specifying a prior
            on the `level_scale` parameter. If `None`, a heuristic default prior is
            constructed based on the provided `observed_time_series`.
            Default value: `None`.
          initial_level_prior: optional `tfd.Distribution` instance specifying a
            prior on the initial level. If `None`, a heuristic default prior is
            constructed based on the provided `observed_time_series`.
            Default value: `None`.
          observed_time_series: optional `float` `Tensor` of shape
            `batch_shape + [T, 1]` (omitting the trailing unit dimension is also
            supported when `T > 1`), specifying an observed time series.
            Any priors not explicitly set will be given default values according to
            the scale of the observed time series (or batch of time series). May
            optionally be an instance of `tfp.sts.MaskedTimeSeries`, which includes
            a mask `Tensor` to specify timesteps with missing observations.
            Default value: `None`.
          name: the name of this model component.
            Default value: 'LocalLevel'.


return self._initial_state_prior<EOL>

Prior distribution on the initial latent state (level and scale).


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

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


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

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


with tf.compat.v1.name_scope(<EOL>name,<EOL>'<STR_LIT>',<EOL>values=[level_scale, slope_mean, slope_scale,<EOL>autoregressive_coef]) as name:<EOL><INDENT>dtype = initial_state_prior.dtype<EOL>level_scale = tf.convert_to_tensor(<EOL>value=level_scale, dtype=dtype, name='<STR_LIT>')<EOL>slope_mean = tf.convert_to_tensor(<EOL>value=slope_mean, dtype=dtype, name='<STR_LIT>')<EOL>slope_scale = tf.convert_to_tensor(<EOL>value=slope_scale, dtype=dtype, name='<STR_LIT>')<EOL>autoregressive_coef = tf.convert_to_tensor(<EOL>value=autoregressive_coef, dtype=dtype, name='<STR_LIT>')<EOL>observation_noise_scale = tf.convert_to_tensor(<EOL>value=observation_noise_scale,<EOL>dtype=dtype,<EOL>name='<STR_LIT>')<EOL>self._level_scale = level_scale<EOL>self._slope_mean = slope_mean<EOL>self._slope_scale = slope_scale<EOL>self._autoregressive_coef = autoregressive_coef<EOL>self._observation_noise_scale = observation_noise_scale<EOL>super(SemiLocalLinearTrendStateSpaceModel, self).__init__(<EOL>num_timesteps=num_timesteps,<EOL>transition_matrix=semilocal_linear_trend_transition_matrix(<EOL>autoregressive_coef),<EOL>transition_noise=semilocal_linear_trend_transition_noise(<EOL>level_scale, slope_mean, slope_scale, autoregressive_coef),<EOL>observation_matrix=tf.constant(<EOL>[[<NUM_LIT:1.>, <NUM_LIT:0.>]], dtype=dtype),<EOL>observation_noise=tfd.MultivariateNormalDiag(<EOL>scale_diag=observation_noise_scale[..., tf.newaxis]),<EOL>initial_state_prior=initial_state_prior,<EOL>initial_step=initial_step,<EOL>allow_nan_stats=allow_nan_stats,<EOL>validate_args=validate_args,<EOL>name=name)<EOL><DEDENT>

Build a state space model implementing a semi-local linear trend.

        Args:
          num_timesteps: Scalar `int` `Tensor` number of timesteps to model
            with this distribution.
          level_scale: Scalar (any additional dimensions are treated as batch
            dimensions) `float` `Tensor` indicating the standard deviation of the
            level transitions.
          slope_mean: Scalar (any additional dimensions are treated as batch
            dimensions) `float` `Tensor` indicating the expected long-term mean of
            the latent slope.
          slope_scale: Scalar (any additional dimensions are treated as batch
            dimensions) `float` `Tensor` indicating the standard deviation of the
            slope transitions.
          autoregressive_coef: Scalar (any additional dimensions are treated as
            batch dimensions) `float` `Tensor` defining the AR1 process on the
            latent slope.
          initial_state_prior: instance of `tfd.MultivariateNormal`
            representing the prior distribution on latent states; must
            have event shape `[2]`.
          observation_noise_scale: Scalar (any additional dimensions are
            treated as batch dimensions) `float` `Tensor` indicating the standard
            deviation of the observation noise.
          initial_step: Optional scalar `int` `Tensor` specifying the starting
            timestep.
            Default value: 0.
          validate_args: Python `bool`. Whether to validate input
            with asserts. If `validate_args` is `False`, and the inputs are
            invalid, correct behavior is not guaranteed.
            Default value: `False`.
          allow_nan_stats: Python `bool`. If `False`, raise an
            exception if a statistic (e.g. mean/mode/etc...) is undefined for any
            batch member. If `True`, batch members with valid parameters leading to
            undefined statistics will return NaN for this statistic.
            Default value: `True`.
          name: Python `str` name prefixed to ops created by this class.
            Default value: "SemiLocalLinearTrendStateSpaceModel".


with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>', values=[observed_time_series]) as name:<EOL><INDENT>if observed_time_series is not None:<EOL><INDENT>_, observed_stddev, observed_initial = sts_util.empirical_statistics(<EOL>observed_time_series)<EOL><DEDENT>else:<EOL><INDENT>observed_stddev, observed_initial = <NUM_LIT:1.>, <NUM_LIT:0.><EOL><DEDENT>if level_scale_prior is None:<EOL><INDENT>level_scale_prior = tfd.LogNormal(<EOL>loc=tf.math.log(<NUM_LIT> * observed_stddev), scale=<NUM_LIT>)<EOL><DEDENT>if slope_mean_prior is None:<EOL><INDENT>slope_mean_prior = tfd.Normal(loc=<NUM_LIT:0.>,<EOL>scale=observed_stddev)<EOL><DEDENT>if slope_scale_prior is None:<EOL><INDENT>slope_scale_prior = tfd.LogNormal(<EOL>loc=tf.math.log(<NUM_LIT> * observed_stddev), scale=<NUM_LIT>)<EOL><DEDENT>if autoregressive_coef_prior is None:<EOL><INDENT>autoregressive_coef_prior = tfd.Normal(<EOL>loc=<NUM_LIT:0.>, scale=tf.ones_like(observed_initial))<EOL><DEDENT>if initial_level_prior is None:<EOL><INDENT>initial_level_prior = tfd.Normal(<EOL>loc=observed_initial,<EOL>scale=tf.abs(observed_initial) + observed_stddev)<EOL><DEDENT>if initial_slope_prior is None:<EOL><INDENT>initial_slope_prior = tfd.Normal(loc=<NUM_LIT:0.>, scale=observed_stddev)<EOL><DEDENT>self._initial_state_prior = tfd.MultivariateNormalDiag(<EOL>loc=tf.stack(<EOL>[initial_level_prior.mean(),<EOL>initial_slope_prior.mean()<EOL>], axis=-<NUM_LIT:1>),<EOL>scale_diag=tf.stack([<EOL>initial_level_prior.stddev(),<EOL>initial_slope_prior.stddev()<EOL>], axis=-<NUM_LIT:1>))<EOL>if constrain_ar_coef_stationary and constrain_ar_coef_positive:<EOL><INDENT>autoregressive_coef_bijector = tfb.Sigmoid()   <EOL><DEDENT>elif constrain_ar_coef_positive:<EOL><INDENT>autoregressive_coef_bijector = tfb.Softplus()  <EOL><DEDENT>elif constrain_ar_coef_stationary:<EOL><INDENT>autoregressive_coef_bijector = tfb.Tanh()      <EOL><DEDENT>else:<EOL><INDENT>autoregressive_coef_bijector = tfb.Identity()  <EOL><DEDENT>super(SemiLocalLinearTrend, self).__init__(<EOL>parameters=[<EOL>Parameter('<STR_LIT>', level_scale_prior, tfb.Softplus()),<EOL>Parameter('<STR_LIT>', slope_mean_prior, tfb.Identity()),<EOL>Parameter('<STR_LIT>', slope_scale_prior, tfb.Softplus()),<EOL>Parameter('<STR_LIT>',<EOL>autoregressive_coef_prior,<EOL>autoregressive_coef_bijector),<EOL>],<EOL>latent_size=<NUM_LIT:2>,<EOL>name=name)<EOL><DEDENT>

Specify a semi-local linear trend model.

        Args:
          level_scale_prior: optional `tfd.Distribution` instance specifying a prior
            on the `level_scale` parameter. If `None`, a heuristic default prior is
            constructed based on the provided `observed_time_series`.
            Default value: `None`.
          slope_mean_prior: optional `tfd.Distribution` instance specifying a prior
            on the `slope_mean` parameter. If `None`, a heuristic default prior is
            constructed based on the provided `observed_time_series`.
            Default value: `None`.
          slope_scale_prior: optional `tfd.Distribution` instance specifying a prior
            on the `slope_scale` parameter. If `None`, a heuristic default prior is
            constructed based on the provided `observed_time_series`.
            Default value: `None`.
          autoregressive_coef_prior: optional `tfd.Distribution` instance specifying
            a prior on the `autoregressive_coef` parameter. If `None`, the default
            prior is a standard `Normal(0., 1.)`. Note that the prior may be
            implicitly truncated by `constrain_ar_coef_stationary` and/or
            `constrain_ar_coef_positive`.
            Default value: `None`.
          initial_level_prior: optional `tfd.Distribution` instance specifying a
            prior on the initial level. If `None`, a heuristic default prior is
            constructed based on the provided `observed_time_series`.
            Default value: `None`.
          initial_slope_prior: optional `tfd.Distribution` instance specifying a
            prior on the initial slope. If `None`, a heuristic default prior is
            constructed based on the provided `observed_time_series`.
            Default value: `None`.
          observed_time_series: optional `float` `Tensor` of shape
            `batch_shape + [T, 1]` (omitting the trailing unit dimension is also
            supported when `T > 1`), specifying an observed time series.
            Any priors not explicitly set will be given default values according to
            the scale of the observed time series (or batch of time series). May
            optionally be an instance of `tfp.sts.MaskedTimeSeries`, which includes
            a mask `Tensor` to specify timesteps with missing observations.
            Default value: `None`.
          constrain_ar_coef_stationary: if `True`, perform inference using a
            parameterization that restricts `autoregressive_coef` to the interval
            `(-1, 1)`, or `(0, 1)` if `force_positive_ar_coef` is also `True`,
            corresponding to stationary processes. This will implicitly truncates
            the support of `autoregressive_coef_prior`.
            Default value: `True`.
          constrain_ar_coef_positive: if `True`, perform inference using a
            parameterization that restricts `autoregressive_coef` to be positive,
            or in `(0, 1)` if `constrain_ar_coef_stationary` is also `True`. This
            will implicitly truncate the support of `autoregressive_coef_prior`.
            Default value: `False`.
          name: the name of this model component.
            Default value: 'SemiLocalLinearTrend'.


return self._initial_state_prior<EOL>


with tf.compat.v1.name_scope(<EOL>name,<EOL>'<STR_LIT>',<EOL>values=[observation_noise_scale, initial_step]) as name:<EOL><INDENT>dtype = tf.debugging.assert_same_float_dtype(component_ssms)<EOL>constant_offset = tf.convert_to_tensor(value=constant_offset,<EOL>name='<STR_LIT>',<EOL>dtype=dtype)<EOL>assertions = []<EOL>if initial_state_prior is None:<EOL><INDENT>initial_state_prior = sts_util.factored_joint_mvn(<EOL>[ssm.initial_state_prior for ssm in component_ssms])<EOL><DEDENT>dtype = initial_state_prior.dtype<EOL>static_num_timesteps = [<EOL>tf.get_static_value(ssm.num_timesteps)<EOL>for ssm in component_ssms<EOL>if tf.get_static_value(ssm.num_timesteps) is not None<EOL>]<EOL>if static_num_timesteps:<EOL><INDENT>num_timesteps = static_num_timesteps[<NUM_LIT:0>]<EOL>if not all([component_timesteps == num_timesteps<EOL>for component_timesteps in static_num_timesteps]):<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'.format(static_num_timesteps))<EOL><DEDENT><DEDENT>else:<EOL><INDENT>num_timesteps = component_ssms[<NUM_LIT:0>].num_timesteps<EOL><DEDENT>if validate_args and len(static_num_timesteps) != len(component_ssms):<EOL><INDENT>assertions += [<EOL>tf.compat.v1.assert_equal(<EOL>num_timesteps,<EOL>ssm.num_timesteps,<EOL>message='<STR_LIT>'<EOL>'<STR_LIT>') for ssm in component_ssms<EOL>]<EOL><DEDENT>def transition_matrix_fn(t):<EOL><INDENT>return tfl.LinearOperatorBlockDiag(<EOL>[ssm.get_transition_matrix_for_timestep(t)<EOL>for ssm in component_ssms])<EOL><DEDENT>def transition_noise_fn(t):<EOL><INDENT>return sts_util.factored_joint_mvn(<EOL>[ssm.get_transition_noise_for_timestep(t)<EOL>for ssm in component_ssms])<EOL><DEDENT>broadcast_batch_shape = tf.convert_to_tensor(<EOL>value=sts_util.broadcast_batch_shape(<EOL>[ssm.get_observation_matrix_for_timestep(initial_step)<EOL>for ssm in component_ssms]), dtype=tf.int32)<EOL>broadcast_obs_matrix = tf.ones(<EOL>tf.concat([broadcast_batch_shape, [<NUM_LIT:1>, <NUM_LIT:1>]], axis=<NUM_LIT:0>), dtype=dtype)<EOL>if assertions:<EOL><INDENT>with tf.control_dependencies(assertions):<EOL><INDENT>broadcast_obs_matrix = tf.identity(broadcast_obs_matrix)<EOL><DEDENT><DEDENT>def observation_matrix_fn(t):<EOL><INDENT>return tfl.LinearOperatorFullMatrix(<EOL>tf.concat([ssm.get_observation_matrix_for_timestep(t).to_dense() *<EOL>broadcast_obs_matrix for ssm in component_ssms],<EOL>axis=-<NUM_LIT:1>))<EOL><DEDENT>offset_vector = constant_offset[..., tf.newaxis]<EOL>if observation_noise_scale is not None:<EOL><INDENT>observation_noise_scale = tf.convert_to_tensor(<EOL>value=observation_noise_scale,<EOL>name='<STR_LIT>',<EOL>dtype=dtype)<EOL>def observation_noise_fn(t):<EOL><INDENT>return tfd.MultivariateNormalDiag(<EOL>loc=sum([ssm.get_observation_noise_for_timestep(t).mean()<EOL>for ssm in component_ssms]) + offset_vector,<EOL>scale_diag=observation_noise_scale[..., tf.newaxis])<EOL><DEDENT><DEDENT>else:<EOL><INDENT>def observation_noise_fn(t):<EOL><INDENT>return sts_util.sum_mvns(<EOL>[tfd.MultivariateNormalDiag(<EOL>loc=offset_vector,<EOL>scale_diag=tf.zeros_like(offset_vector))] +<EOL>[ssm.get_observation_noise_for_timestep(t)<EOL>for ssm in component_ssms])<EOL><DEDENT><DEDENT>super(AdditiveStateSpaceModel, self).__init__(<EOL>num_timesteps=num_timesteps,<EOL>transition_matrix=transition_matrix_fn,<EOL>transition_noise=transition_noise_fn,<EOL>observation_matrix=observation_matrix_fn,<EOL>observation_noise=observation_noise_fn,<EOL>initial_state_prior=initial_state_prior,<EOL>initial_step=initial_step,<EOL>validate_args=validate_args,<EOL>allow_nan_stats=allow_nan_stats,<EOL>name=name)<EOL><DEDENT>

Build a state space model representing the sum of component models.

        Args:
          component_ssms: Python `list` containing one or more
            `tfd.LinearGaussianStateSpaceModel` instances. The components
            will in general implement different time-series models, with possibly
            different `latent_size`, but they must have the same `dtype`, event
            shape (`num_timesteps` and `observation_size`), and their batch shapes
            must broadcast to a compatible batch shape.
          constant_offset: scalar `float` `Tensor`, or batch of scalars,
            specifying a constant value added to the sum of outputs from the
            component models. This allows the components to model the shifted series
            `observed_time_series - constant_offset`.
            Default value: `0.`
          observation_noise_scale: Optional scalar `float` `Tensor` indicating the
            standard deviation of the observation noise. May contain additional
            batch dimensions, which must broadcast with the batch shape of elements
            in `component_ssms`. If `observation_noise_scale` is specified for the
            `AdditiveStateSpaceModel`, the observation noise scales of component
            models are ignored. If `None`, the observation noise scale is derived
            by summing the noise variances of the component models, i.e.,
            `observation_noise_scale = sqrt(sum(
            [ssm.observation_noise_scale**2 for ssm in component_ssms]))`.
          initial_state_prior: Optional instance of `tfd.MultivariateNormal`
            representing a prior distribution on the latent state at time
            `initial_step`. If `None`, defaults to the independent priors from
            component models, i.e.,
            `[component.initial_state_prior for component in component_ssms]`.
            Default value: `None`.
          initial_step: Optional scalar `int` `Tensor` specifying the starting
            timestep.
            Default value: 0.
          validate_args: Python `bool`. Whether to validate input
            with asserts. If `validate_args` is `False`, and the inputs are
            invalid, correct behavior is not guaranteed.
            Default value: `False`.
          allow_nan_stats: Python `bool`. If `False`, raise an
            exception if a statistic (e.g. mean/mode/etc...) is undefined for any
            batch member. If `True`, batch members with valid parameters leading to
            undefined statistics will return NaN for this statistic.
            Default value: `True`.
          name: Python `str` name prefixed to ops created by this class.
            Default value: "AdditiveStateSpaceModel".

        Raises:
          ValueError: if components have different `num_timesteps`.


with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>', values=[observed_time_series]) as name:<EOL><INDENT>if observed_time_series is not None:<EOL><INDENT>observed_mean, observed_stddev, _ = (<EOL>sts_util.empirical_statistics(observed_time_series))<EOL><DEDENT>else:<EOL><INDENT>observed_mean, observed_stddev = <NUM_LIT:0.>, <NUM_LIT:1.><EOL><DEDENT>if observation_noise_scale_prior is None:<EOL><INDENT>observation_noise_scale_prior = tfd.LogNormal(<EOL>loc=tf.math.log(<NUM_LIT> * observed_stddev), scale=<NUM_LIT>)<EOL><DEDENT>if constant_offset is None:<EOL><INDENT>constant_offset = observed_mean<EOL><DEDENT>component_names = [c.name for c in components]<EOL>if len(component_names) != len(set(component_names)):<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'.format(component_names))<EOL><DEDENT>components_by_name = collections.OrderedDict(<EOL>[(c.name, c) for c in components])<EOL>parameters = [<EOL>Parameter('<STR_LIT>', observation_noise_scale_prior,<EOL>tfb.Softplus()),<EOL>] + [Parameter(name='<STR_LIT>'.format(component.name, parameter.name),<EOL>prior=parameter.prior,<EOL>bijector=parameter.bijector)<EOL>for component in components for parameter in component.parameters]<EOL>self._components = components<EOL>self._components_by_name = components_by_name<EOL>self._constant_offset = constant_offset<EOL>super(Sum, self).__init__(<EOL>parameters=parameters,<EOL>latent_size=sum(<EOL>[component.latent_size for component in components]),<EOL>name=name)<EOL><DEDENT>

Specify a structural time series model representing a sum of components.

        Args:
          components: Python `list` of one or more StructuralTimeSeries instances.
            These must have unique names.
          constant_offset: optional scalar `float` `Tensor`, or batch of scalars,
            specifying a constant value added to the sum of outputs from the
            component models. This allows the components to model the shifted series
            `observed_time_series - constant_offset`. If `None`, this is set to the
            mean of the provided `observed_time_series`.
            Default value: `None`.
          observation_noise_scale_prior: optional `tfd.Distribution` instance
            specifying a prior on `observation_noise_scale`. If `None`, a heuristic
            default prior is constructed based on the provided
            `observed_time_series`.
            Default value: `None`.
          observed_time_series: optional `float` `Tensor` of shape
            `batch_shape + [T, 1]` (omitting the trailing unit dimension is also
            supported when `T > 1`), specifying an observed time series. This is
            used to set the constant offset, if not provided, and to construct a
            default heuristic `observation_noise_scale_prior` if not provided. May
            optionally be an instance of `tfp.sts.MaskedTimeSeries`, which includes
            a mask `Tensor` to specify timesteps with missing observations.
            Default value: `None`.
          name: Python `str` name of this model component; used as `name_scope`
            for ops created by this class.
            Default value: 'Sum'.

        Raises:
          ValueError: if components do not have unique names.


return self._components<EOL>

List of component `StructuralTimeSeries` models.


return self._components_by_name<EOL>

OrderedDict mapping component names to components.


return self._constant_offset<EOL>

Constant value subtracted from observed data.


with tf.compat.v1.name_scope('<STR_LIT>'):<EOL><INDENT>param_map = self._canonicalize_param_vals_as_map(param_vals)<EOL>param_vals_list = [param_map[p.name] for p in self.parameters]<EOL>remaining_param_vals = param_vals_list[<NUM_LIT:1>:]<EOL>component_ssms = []<EOL>for component in self.components:<EOL><INDENT>num_parameters = len(component.parameters)<EOL>component_param_vals = remaining_param_vals[:num_parameters]<EOL>remaining_param_vals = remaining_param_vals[num_parameters:]<EOL>component_ssms.append(<EOL>component.make_state_space_model(<EOL>num_timesteps,<EOL>param_vals=component_param_vals,<EOL>initial_step=initial_step))<EOL><DEDENT><DEDENT>return component_ssms<EOL>

Build an ordered list of Distribution instances for component models.

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

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


num_steps_per_cycle = np.sum(num_steps_per_season)<EOL>changepoints = np.cumsum(np.ravel(num_steps_per_season)) - <NUM_LIT:1><EOL>def is_last_day_of_season(t):<EOL><INDENT>t_ = dist_util.maybe_get_static_value(t)<EOL>if t_ is not None:  <EOL><INDENT>step_in_cycle = t_ % num_steps_per_cycle<EOL>return any(step_in_cycle == changepoints)<EOL><DEDENT>else:<EOL><INDENT>step_in_cycle = tf.math.floormod(t, num_steps_per_cycle)<EOL>return tf.reduce_any(<EOL>input_tensor=tf.equal(step_in_cycle, changepoints))<EOL><DEDENT><DEDENT>return is_last_day_of_season<EOL>

Build utility method to compute whether the season is changing.


<EOL>effects_to_residuals_fullrank = np.eye(num_seasons) - <NUM_LIT:1.>/num_seasons<EOL>effects_to_residuals_fullrank[-<NUM_LIT:1>, :] = <NUM_LIT:1.>/num_seasons  <EOL>residuals_to_effects_fullrank = np.linalg.inv(effects_to_residuals_fullrank)<EOL>effects_to_residuals = effects_to_residuals_fullrank[:-<NUM_LIT:1>, :]<EOL>residuals_to_effects = residuals_to_effects_fullrank[:, :-<NUM_LIT:1>]<EOL>effects_to_residuals = tf.cast(<EOL>effects_to_residuals, dtype=dtype, name='<STR_LIT>')<EOL>residuals_to_effects = tf.cast(<EOL>residuals_to_effects, dtype=dtype, name='<STR_LIT>')<EOL>return effects_to_residuals, residuals_to_effects<EOL>

Build change-of-basis matrices for constrained seasonal effects.

    This method builds the matrix that transforms seasonal effects into
    effect residuals (differences from the mean effect), and additionally
    projects these residuals onto the subspace where the mean effect is zero.

    See `ConstrainedSeasonalStateSpaceModel` for mathematical details.

    Args:
      num_seasons: scalar `int` number of seasons.
      dtype: TensorFlow `dtype` for the returned values.
    Returns:
      effects_to_residuals: `Tensor` of shape
        `[num_seasons-1, num_seasons]`, such that `differences_from_mean_effect =
        matmul(effects_to_residuals, seasonal_effects)`.  In the
        notation of `ConstrainedSeasonalStateSpaceModel`, this is
        `effects_to_residuals = P * R`.
      residuals_to_effects: the (pseudo)-inverse of the above; a
        `Tensor` of shape `[num_seasons, num_seasons-1]`. In the
        notation of `ConstrainedSeasonalStateSpaceModel`, this is
        `residuals_to_effects = R^{-1} * P'`.


def build_seasonal_transition_matrix(<EOL>num_seasons, is_last_day_of_season, dtype,<EOL>basis_change_matrix=None, basis_change_matrix_inv=None):

with tf.compat.v1.name_scope('<STR_LIT>'):<EOL><INDENT>seasonal_permutation = np.concatenate(<EOL>[np.arange(<NUM_LIT:1>, num_seasons), [<NUM_LIT:0>]], axis=<NUM_LIT:0>)<EOL>seasonal_permutation_matrix = tf.constant(<EOL>np.eye(num_seasons)[seasonal_permutation], dtype=dtype)<EOL>if basis_change_matrix is not None:<EOL><INDENT>seasonal_permutation_matrix = tf.matmul(<EOL>basis_change_matrix,<EOL>tf.matmul(seasonal_permutation_matrix, basis_change_matrix_inv))<EOL><DEDENT>identity_matrix = tf.eye(<EOL>tf.shape(input=seasonal_permutation_matrix)[-<NUM_LIT:1>], dtype=dtype)<EOL>def seasonal_transition_matrix(t):<EOL><INDENT>return tf.linalg.LinearOperatorFullMatrix(<EOL>matrix=dist_util.pick_scalar_condition(<EOL>is_last_day_of_season(t),<EOL>seasonal_permutation_matrix,<EOL>identity_matrix))<EOL><DEDENT><DEDENT>return seasonal_transition_matrix<EOL>

Build a function computing transitions for a seasonal effect model.


def build_seasonal_transition_noise(<EOL>drift_scale, num_seasons, is_last_day_of_season):

<EOL>drift_scale_diag = tf.stack(<EOL>[tf.zeros_like(drift_scale)] * (num_seasons - <NUM_LIT:1>) + [drift_scale],<EOL>axis=-<NUM_LIT:1>)<EOL>def seasonal_transition_noise(t):<EOL><INDENT>noise_scale_diag = dist_util.pick_scalar_condition(<EOL>is_last_day_of_season(t),<EOL>drift_scale_diag,<EOL>tf.zeros_like(drift_scale_diag))<EOL>return tfd.MultivariateNormalDiag(<EOL>loc=tf.zeros(num_seasons, dtype=drift_scale.dtype),<EOL>scale_diag=noise_scale_diag)<EOL><DEDENT>return seasonal_transition_noise<EOL>

Build the transition noise model for a SeasonalStateSpaceModel.


def build_constrained_seasonal_transition_noise(<EOL>drift_scale, num_seasons, is_last_day_of_season):

<EOL>drift_scale_tril_nonzeros = tf.concat([<EOL>tf.ones([num_seasons - <NUM_LIT:1>, <NUM_LIT:1>], dtype=drift_scale.dtype),<EOL>tf.zeros([num_seasons - <NUM_LIT:1>, num_seasons - <NUM_LIT:2>], dtype=drift_scale.dtype)],<EOL>axis=-<NUM_LIT:1>)<EOL>drift_scale_tril = (drift_scale_tril_nonzeros *<EOL>drift_scale[..., tf.newaxis, tf.newaxis] / num_seasons)<EOL>def seasonal_transition_noise(t):<EOL><INDENT>noise_scale_tril = dist_util.pick_scalar_condition(<EOL>is_last_day_of_season(t),<EOL>drift_scale_tril,<EOL>tf.zeros_like(drift_scale_tril))<EOL>return tfd.MultivariateNormalTriL(<EOL>loc=tf.zeros(num_seasons-<NUM_LIT:1>, dtype=drift_scale.dtype),<EOL>scale_tril=noise_scale_tril)<EOL><DEDENT>return seasonal_transition_noise<EOL>

Build transition noise distribution for a ConstrainedSeasonalSSM.


with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>',<EOL>values=[drift_scale, observation_noise_scale]) as name:<EOL><INDENT>dtype = initial_state_prior.dtype<EOL>drift_scale = tf.convert_to_tensor(<EOL>value=drift_scale, name='<STR_LIT>', dtype=dtype)<EOL>observation_noise_scale = tf.convert_to_tensor(<EOL>value=observation_noise_scale,<EOL>name='<STR_LIT>',<EOL>dtype=dtype)<EOL>num_steps_per_season = np.squeeze(np.asarray(num_steps_per_season))<EOL>if num_steps_per_season.ndim == <NUM_LIT:0>:  <EOL><INDENT>num_steps_per_season = np.tile(num_steps_per_season, num_seasons)<EOL><DEDENT>elif ((num_steps_per_season.ndim <= <NUM_LIT:2>)  <EOL>and (num_steps_per_season.shape[-<NUM_LIT:1>] != num_seasons)):<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'.format(<EOL>num_seasons, num_steps_per_season.shape))<EOL><DEDENT>is_last_day_of_season = build_is_last_day_of_season(num_steps_per_season)<EOL>seasonal_transition_matrix = build_seasonal_transition_matrix(<EOL>num_seasons=num_seasons,<EOL>is_last_day_of_season=is_last_day_of_season,<EOL>dtype=dtype)<EOL>seasonal_transition_noise = build_seasonal_transition_noise(<EOL>drift_scale, num_seasons, is_last_day_of_season)<EOL>observation_matrix = tf.concat([<EOL>tf.ones([<NUM_LIT:1>, <NUM_LIT:1>], dtype=dtype),<EOL>tf.zeros([<NUM_LIT:1>, num_seasons-<NUM_LIT:1>], dtype=dtype)], axis=-<NUM_LIT:1>)<EOL>self._drift_scale = drift_scale<EOL>self._observation_noise_scale = observation_noise_scale<EOL>self._num_seasons = num_seasons<EOL>self._num_steps_per_season = num_steps_per_season<EOL>super(SeasonalStateSpaceModel, self).__init__(<EOL>num_timesteps=num_timesteps,<EOL>transition_matrix=seasonal_transition_matrix,<EOL>transition_noise=seasonal_transition_noise,<EOL>observation_matrix=observation_matrix,<EOL>observation_noise=tfd.MultivariateNormalDiag(<EOL>scale_diag=observation_noise_scale[..., tf.newaxis]),<EOL>initial_state_prior=initial_state_prior,<EOL>initial_step=initial_step,<EOL>allow_nan_stats=allow_nan_stats,<EOL>validate_args=validate_args,<EOL>name=name)<EOL><DEDENT>

Build a seasonal effect state space model.

        {seasonal_init_args}


return self._drift_scale<EOL>

Standard deviation of the drift in effects between seasonal cycles.


Standard deviation of the observation noise.


return self._num_seasons<EOL>


return self._num_steps_per_season<EOL>


with tf.compat.v1.name_scope(<EOL>name, '<STR_LIT>',<EOL>values=[drift_scale, observation_noise_scale]) as name:<EOL><INDENT>dtype = initial_state_prior.dtype<EOL>drift_scale = tf.convert_to_tensor(<EOL>value=drift_scale, name='<STR_LIT>', dtype=dtype)<EOL>observation_noise_scale = tf.convert_to_tensor(<EOL>value=observation_noise_scale,<EOL>name='<STR_LIT>',<EOL>dtype=dtype)<EOL>num_steps_per_season = np.squeeze(np.asarray(num_steps_per_season))<EOL>if num_steps_per_season.ndim == <NUM_LIT:0>:  <EOL><INDENT>num_steps_per_season = np.tile(num_steps_per_season, num_seasons)<EOL><DEDENT>elif ((num_steps_per_season.ndim <= <NUM_LIT:2>)  <EOL>and (num_steps_per_season.shape[-<NUM_LIT:1>] != num_seasons)):<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'.format(<EOL>num_seasons, num_steps_per_season.shape))<EOL><DEDENT>is_last_day_of_season = build_is_last_day_of_season(num_steps_per_season)<EOL>[<EOL>effects_to_residuals,<EOL>residuals_to_effects<EOL>] = build_effects_to_residuals_matrix(num_seasons, dtype=dtype)<EOL>seasonal_transition_matrix = build_seasonal_transition_matrix(<EOL>num_seasons=num_seasons,<EOL>is_last_day_of_season=is_last_day_of_season,<EOL>dtype=dtype,<EOL>basis_change_matrix=effects_to_residuals,<EOL>basis_change_matrix_inv=residuals_to_effects)<EOL>seasonal_transition_noise = build_constrained_seasonal_transition_noise(<EOL>drift_scale, num_seasons, is_last_day_of_season)<EOL>observation_matrix = tf.concat(<EOL>[tf.ones([<NUM_LIT:1>, <NUM_LIT:1>], dtype=dtype),<EOL>tf.zeros([<NUM_LIT:1>, num_seasons-<NUM_LIT:1>], dtype=dtype)], axis=-<NUM_LIT:1>)<EOL>observation_matrix = tf.matmul(observation_matrix, residuals_to_effects)<EOL>self._drift_scale = drift_scale<EOL>self._observation_noise_scale = observation_noise_scale<EOL>self._num_seasons = num_seasons<EOL>self._num_steps_per_season = num_steps_per_season<EOL>super(ConstrainedSeasonalStateSpaceModel, self).__init__(<EOL>num_timesteps=num_timesteps,<EOL>transition_matrix=seasonal_transition_matrix,<EOL>transition_noise=seasonal_transition_noise,<EOL>observation_matrix=observation_matrix,<EOL>observation_noise=tfd.MultivariateNormalDiag(<EOL>scale_diag=observation_noise_scale[..., tf.newaxis]),<EOL>initial_state_prior=initial_state_prior,<EOL>initial_step=initial_step,<EOL>allow_nan_stats=allow_nan_stats,<EOL>validate_args=validate_args,<EOL>name=name)<EOL><DEDENT>

Build a seasonal effect state space model with a zero-sum constraint.

        {seasonal_init_args}


return self._drift_scale<EOL>