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

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

    with tf.name_scope("mean_joint"):

      # The initial timestep is a special case, since we sample the
      # latent state from the prior rather than the transition model.

      with tf.control_dependencies(self.runtime_assertions):
        # Broadcast to ensure we represent the full batch shape.
        initial_latent_mean = _broadcast_to_shape(
            self.initial_state_prior.mean()[..., tf.newaxis],
            tf.concat([self.batch_shape_tensor(),
                       [self.latent_size, 1]], axis=0))

      initial_observation_mean = _propagate_mean(
          initial_latent_mean,
          self.get_observation_matrix_for_timestep(self.initial_step),
          self.get_observation_noise_for_timestep(self.initial_step))

      mean_step = build_kalman_mean_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_means, observation_means) = tf.scan(
          mean_step,
          elems=tf.range(self.initial_step+1, self.final_step),
          initializer=(initial_latent_mean, initial_observation_mean))

      # Squish the initial step back on top of the other (scanned) timesteps
      latent_means = tf.concat([initial_latent_mean[tf.newaxis, ...],
                                latent_means], axis=0)
      observation_means = tf.concat([initial_observation_mean[tf.newaxis, ...],
                                     observation_means], axis=0)

      # Put dimensions back in order. The samples we've computed have
      # shape `[num_timesteps, batch_shape, size, 1]`, where `size`
      # is the dimension of the latent or observation spaces
      # respectively, but we want to return values with shape
      # `[batch_shape, num_timesteps, size]`.
      latent_means = tf.squeeze(latent_means, -1)
      latent_means = distribution_util.move_dimension(latent_means, 0, -2)
      observation_means = tf.squeeze(observation_means, -1)
      observation_means = distribution_util.move_dimension(
          observation_means, 0, -2)

      return latent_means, observation_means
def _joint_mean(self):
    with tf.name_scope("mean_joint"):

      # The initial timestep is a special case, since we sample the
      # latent state from the prior rather than the transition model.

      with tf.control_dependencies(self.runtime_assertions):
        # Broadcast to ensure we represent the full batch shape.
        initial_latent_mean = _broadcast_to_shape(
            self.initial_state_prior.mean()[..., tf.newaxis],
            tf.concat([self.batch_shape_tensor(),
                       [self.latent_size, 1]], axis=0))

      initial_observation_mean = _propagate_mean(
          initial_latent_mean,
          self.get_observation_matrix_for_timestep(self.initial_step),
          self.get_observation_noise_for_timestep(self.initial_step))

      mean_step = build_kalman_mean_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_means, observation_means) = tf.scan(
          mean_step,
          elems=tf.range(self.initial_step+1, self.final_step),
          initializer=(initial_latent_mean, initial_observation_mean))

      # Squish the initial step back on top of the other (scanned) timesteps
      latent_means = tf.concat([initial_latent_mean[tf.newaxis, ...],
                                latent_means], axis=0)
      observation_means = tf.concat([initial_observation_mean[tf.newaxis, ...],
                                     observation_means], axis=0)

      # Put dimensions back in order. The samples we've computed have
      # shape `[num_timesteps, batch_shape, size, 1]`, where `size`
      # is the dimension of the latent or observation spaces
      # respectively, but we want to return values with shape
      # `[batch_shape, num_timesteps, size]`.
      latent_means = tf.squeeze(latent_means, -1)
      latent_means = distribution_util.move_dimension(latent_means, 0, -2)
      observation_means = tf.squeeze(observation_means, -1)
      observation_means = distribution_util.move_dimension(
          observation_means, 0, -2)

      return latent_means, observation_means
def _joint_mean ( self ) : with tf . name_scope ( "mean_joint" ) : # The initial timestep is a special case, since we sample the # latent state from the prior rather than the transition model. with tf . control_dependencies ( self . runtime_assertions ) : # Broadcast to ensure we represent the full batch shape. initial_latent_mean = _broadcast_to_shape ( self . initial_state_prior . mean ( ) [ ... , tf . newaxis ] , tf . concat ( [ self . batch_shape_tensor ( ) , [ self . latent_size , 1 ] ] , axis = 0 ) ) initial_observation_mean = _propagate_mean ( initial_latent_mean , self . get_observation_matrix_for_timestep ( self . initial_step ) , self . get_observation_noise_for_timestep ( self . initial_step ) ) mean_step = build_kalman_mean_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_means , observation_means ) = tf . scan ( mean_step , elems = tf . range ( self . initial_step + 1 , self . final_step ) , initializer = ( initial_latent_mean , initial_observation_mean ) ) # Squish the initial step back on top of the other (scanned) timesteps latent_means = tf . concat ( [ initial_latent_mean [ tf . newaxis , ... ] , latent_means ] , axis = 0 ) observation_means = tf . concat ( [ initial_observation_mean [ tf . newaxis , ... ] , observation_means ] , axis = 0 ) # Put dimensions back in order. The samples we've computed have # shape `[num_timesteps, batch_shape, size, 1]`, where `size` # is the dimension of the latent or observation spaces # respectively, but we want to return values with shape # `[batch_shape, num_timesteps, size]`. latent_means = tf . squeeze ( latent_means , - 1 ) latent_means = distribution_util . move_dimension ( latent_means , 0 , - 2 ) observation_means = tf . squeeze ( observation_means , - 1 ) observation_means = distribution_util . move_dimension ( observation_means , 0 , - 2 ) return latent_means , observation_means
Compute prior means for all variables via dynamic programming.

    Returns:
      latent_means: Prior means of latent states `z_t`, as a `Tensor`
        of shape `batch_shape + [num_timesteps, latent_size]`
      observation_means: Prior covariance matrices of observations
        `x_t`, as a `Tensor` of shape `batch_shape + [num_timesteps,
        observation_size]`
Compute prior means for all variables via dynamic programming .