tensorflow_probability/python/distributions/linear_gaussian_ssm.py
def _joint_covariances(self):
    """Compute prior covariances for all variables via dynamic programming.

    Returns:
      latent_covs: Prior covariance matrices of latent states `z_t`, as
        a `Tensor` of shape `batch_shape + [num_timesteps,
        latent_size, latent_size]`
      observation_covs: Prior covariance matrices of observations
        `x_t`, as a `Tensor` of shape `batch_shape + [num_timesteps,
        observation_size, observation_size]`
    """

    with tf.name_scope("covariance_joint"):

      with tf.control_dependencies(self.runtime_assertions):
        initial_latent_cov = _broadcast_to_shape(
            self.initial_state_prior.covariance(),
            tf.concat([self.batch_shape_tensor(),
                       [self.latent_size, self.latent_size]], axis=0))

      initial_observation_cov = _propagate_cov(
          initial_latent_cov,
          self.get_observation_matrix_for_timestep(self.initial_step),
          self.get_observation_noise_for_timestep(self.initial_step))

      cov_step = build_kalman_cov_step(
          self.get_transition_matrix_for_timestep,
          self.get_transition_noise_for_timestep,
          self.get_observation_matrix_for_timestep,
          self.get_observation_noise_for_timestep)

      # Scan over all timesteps following the initial step.
      (latent_covs, observation_covs) = tf.scan(
          cov_step,
          elems=tf.range(self.initial_step+1, self.final_step),
          initializer=(initial_latent_cov, initial_observation_cov))

      # Squish the initial step back on top of the other (scanned) timesteps
      latent_covs = tf.concat([initial_latent_cov[tf.newaxis, ...],
                               latent_covs], axis=0)
      observation_covs = tf.concat([initial_observation_cov[tf.newaxis, ...],
                                    observation_covs], axis=0)

      # Put dimensions back in order. The samples we've computed have
      # shape `[num_timesteps, batch_shape, size, size]`, where `size`
      # is the dimension of the state or observation spaces
      # respectively, but we want to return values with shape
      # `[batch_shape, num_timesteps, size, size]`.
      latent_covs = distribution_util.move_dimension(latent_covs, 0, -3)
      observation_covs = distribution_util.move_dimension(
          observation_covs, 0, -3)
      return latent_covs, observation_covs
def _joint_covariances(self):
    """Compute prior covariances for all variables via dynamic programming.

    Returns:
      latent_covs: Prior covariance matrices of latent states `z_t`, as
        a `Tensor` of shape `batch_shape + [num_timesteps,
        latent_size, latent_size]`
      observation_covs: Prior covariance matrices of observations
        `x_t`, as a `Tensor` of shape `batch_shape + [num_timesteps,
        observation_size, observation_size]`
    """

    with tf.name_scope("covariance_joint"):

      with tf.control_dependencies(self.runtime_assertions):
        initial_latent_cov = _broadcast_to_shape(
            self.initial_state_prior.covariance(),
            tf.concat([self.batch_shape_tensor(),
                       [self.latent_size, self.latent_size]], axis=0))

      initial_observation_cov = _propagate_cov(
          initial_latent_cov,
          self.get_observation_matrix_for_timestep(self.initial_step),
          self.get_observation_noise_for_timestep(self.initial_step))

      cov_step = build_kalman_cov_step(
          self.get_transition_matrix_for_timestep,
          self.get_transition_noise_for_timestep,
          self.get_observation_matrix_for_timestep,
          self.get_observation_noise_for_timestep)

      # Scan over all timesteps following the initial step.
      (latent_covs, observation_covs) = tf.scan(
          cov_step,
          elems=tf.range(self.initial_step+1, self.final_step),
          initializer=(initial_latent_cov, initial_observation_cov))

      # Squish the initial step back on top of the other (scanned) timesteps
      latent_covs = tf.concat([initial_latent_cov[tf.newaxis, ...],
                               latent_covs], axis=0)
      observation_covs = tf.concat([initial_observation_cov[tf.newaxis, ...],
                                    observation_covs], axis=0)

      # Put dimensions back in order. The samples we've computed have
      # shape `[num_timesteps, batch_shape, size, size]`, where `size`
      # is the dimension of the state or observation spaces
      # respectively, but we want to return values with shape
      # `[batch_shape, num_timesteps, size, size]`.
      latent_covs = distribution_util.move_dimension(latent_covs, 0, -3)
      observation_covs = distribution_util.move_dimension(
          observation_covs, 0, -3)
      return latent_covs, observation_covs
Compute prior covariances for all variables via dynamic programming.

    Returns:
      latent_covs: Prior covariance matrices of latent states `z_t`, as
        a `Tensor` of shape `batch_shape + [num_timesteps,
        latent_size, latent_size]`
      observation_covs: Prior covariance matrices of observations
        `x_t`, as a `Tensor` of shape `batch_shape + [num_timesteps,
        observation_size, observation_size]`