tensorflow_probability/python/distributions/linear_gaussian_ssm.py
def _joint_sample_n(self, n, seed=None):
    """Draw a joint sample from the prior over latents and observations."""

    with tf.name_scope("sample_n_joint"):
      stream = seed_stream.SeedStream(
          seed, salt="LinearGaussianStateSpaceModel_sample_n_joint")

      sample_and_batch_shape = distribution_util.prefer_static_value(
          tf.concat([[n], self.batch_shape_tensor()],
                    axis=0))

      # Sample the initial timestep from the prior.  Since we want
      # this sample to have full batch shape (not just the batch shape
      # of the self.initial_state_prior object which might in general be
      # smaller), we augment the sample shape to include whatever
      # extra batch dimensions are required.
      with tf.control_dependencies(self.runtime_assertions):
        initial_latent = self.initial_state_prior.sample(
            sample_shape=_augment_sample_shape(
                self.initial_state_prior,
                sample_and_batch_shape,
                self.validate_args),
            seed=stream())

        # Add a dummy dimension so that matmul() does matrix-vector
        # multiplication.
        initial_latent = initial_latent[..., tf.newaxis]

      initial_observation_matrix = (
          self.get_observation_matrix_for_timestep(self.initial_step))
      initial_observation_noise = (
          self.get_observation_noise_for_timestep(self.initial_step))

      initial_observation_pred = initial_observation_matrix.matmul(
          initial_latent)
      initial_observation = (initial_observation_pred +
                             initial_observation_noise.sample(
                                 sample_shape=_augment_sample_shape(
                                     initial_observation_noise,
                                     sample_and_batch_shape,
                                     self.validate_args),
                                 seed=stream())[..., tf.newaxis])

      sample_step = build_kalman_sample_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,
          full_sample_and_batch_shape=sample_and_batch_shape,
          stream=stream,
          validate_args=self.validate_args)

      # Scan over all timesteps to sample latents and observations.
      (latents, observations) = tf.scan(
          sample_step,
          elems=tf.range(self.initial_step+1, self.final_step),
          initializer=(initial_latent, initial_observation))

      # Combine the initial sampled timestep with the remaining timesteps.
      latents = tf.concat([initial_latent[tf.newaxis, ...],
                           latents], axis=0)
      observations = tf.concat([initial_observation[tf.newaxis, ...],
                                observations], axis=0)

      # Put dimensions back in order. The samples we've computed are
      # ordered by timestep, with shape `[num_timesteps, num_samples,
      # batch_shape, size, 1]` where `size` represents `latent_size`
      # or `observation_size` respectively. But timesteps are really
      # part of each probabilistic event, so we need to return a Tensor
      # of shape `[num_samples, batch_shape, num_timesteps, size]`.
      latents = tf.squeeze(latents, -1)
      latents = distribution_util.move_dimension(latents, 0, -2)
      observations = tf.squeeze(observations, -1)
      observations = distribution_util.move_dimension(observations, 0, -2)

    return latents, observations
def _joint_sample_n(self, n, seed=None):
    with tf.name_scope("sample_n_joint"):
      stream = seed_stream.SeedStream(
          seed, salt="LinearGaussianStateSpaceModel_sample_n_joint")

      sample_and_batch_shape = distribution_util.prefer_static_value(
          tf.concat([[n], self.batch_shape_tensor()],
                    axis=0))

      # Sample the initial timestep from the prior.  Since we want
      # this sample to have full batch shape (not just the batch shape
      # of the self.initial_state_prior object which might in general be
      # smaller), we augment the sample shape to include whatever
      # extra batch dimensions are required.
      with tf.control_dependencies(self.runtime_assertions):
        initial_latent = self.initial_state_prior.sample(
            sample_shape=_augment_sample_shape(
                self.initial_state_prior,
                sample_and_batch_shape,
                self.validate_args),
            seed=stream())

        # Add a dummy dimension so that matmul() does matrix-vector
        # multiplication.
        initial_latent = initial_latent[..., tf.newaxis]

      initial_observation_matrix = (
          self.get_observation_matrix_for_timestep(self.initial_step))
      initial_observation_noise = (
          self.get_observation_noise_for_timestep(self.initial_step))

      initial_observation_pred = initial_observation_matrix.matmul(
          initial_latent)
      initial_observation = (initial_observation_pred +
                             initial_observation_noise.sample(
                                 sample_shape=_augment_sample_shape(
                                     initial_observation_noise,
                                     sample_and_batch_shape,
                                     self.validate_args),
                                 seed=stream())[..., tf.newaxis])

      sample_step = build_kalman_sample_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,
          full_sample_and_batch_shape=sample_and_batch_shape,
          stream=stream,
          validate_args=self.validate_args)

      # Scan over all timesteps to sample latents and observations.
      (latents, observations) = tf.scan(
          sample_step,
          elems=tf.range(self.initial_step+1, self.final_step),
          initializer=(initial_latent, initial_observation))

      # Combine the initial sampled timestep with the remaining timesteps.
      latents = tf.concat([initial_latent[tf.newaxis, ...],
                           latents], axis=0)
      observations = tf.concat([initial_observation[tf.newaxis, ...],
                                observations], axis=0)

      # Put dimensions back in order. The samples we've computed are
      # ordered by timestep, with shape `[num_timesteps, num_samples,
      # batch_shape, size, 1]` where `size` represents `latent_size`
      # or `observation_size` respectively. But timesteps are really
      # part of each probabilistic event, so we need to return a Tensor
      # of shape `[num_samples, batch_shape, num_timesteps, size]`.
      latents = tf.squeeze(latents, -1)
      latents = distribution_util.move_dimension(latents, 0, -2)
      observations = tf.squeeze(observations, -1)
      observations = distribution_util.move_dimension(observations, 0, -2)

    return latents, observations
def _joint_sample_n ( self , n , seed = None ) : with tf . name_scope ( "sample_n_joint" ) : stream = seed_stream . SeedStream ( seed , salt = "LinearGaussianStateSpaceModel_sample_n_joint" ) sample_and_batch_shape = distribution_util . prefer_static_value ( tf . concat ( [ [ n ] , self . batch_shape_tensor ( ) ] , axis = 0 ) ) # Sample the initial timestep from the prior.  Since we want # this sample to have full batch shape (not just the batch shape # of the self.initial_state_prior object which might in general be # smaller), we augment the sample shape to include whatever # extra batch dimensions are required. with tf . control_dependencies ( self . runtime_assertions ) : initial_latent = self . initial_state_prior . sample ( sample_shape = _augment_sample_shape ( self . initial_state_prior , sample_and_batch_shape , self . validate_args ) , seed = stream ( ) ) # Add a dummy dimension so that matmul() does matrix-vector # multiplication. initial_latent = initial_latent [ ... , tf . newaxis ] initial_observation_matrix = ( self . get_observation_matrix_for_timestep ( self . initial_step ) ) initial_observation_noise = ( self . get_observation_noise_for_timestep ( self . initial_step ) ) initial_observation_pred = initial_observation_matrix . matmul ( initial_latent ) initial_observation = ( initial_observation_pred + initial_observation_noise . sample ( sample_shape = _augment_sample_shape ( initial_observation_noise , sample_and_batch_shape , self . validate_args ) , seed = stream ( ) ) [ ... , tf . newaxis ] ) sample_step = build_kalman_sample_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 , full_sample_and_batch_shape = sample_and_batch_shape , stream = stream , validate_args = self . validate_args ) # Scan over all timesteps to sample latents and observations. ( latents , observations ) = tf . scan ( sample_step , elems = tf . range ( self . initial_step + 1 , self . final_step ) , initializer = ( initial_latent , initial_observation ) ) # Combine the initial sampled timestep with the remaining timesteps. latents = tf . concat ( [ initial_latent [ tf . newaxis , ... ] , latents ] , axis = 0 ) observations = tf . concat ( [ initial_observation [ tf . newaxis , ... ] , observations ] , axis = 0 ) # Put dimensions back in order. The samples we've computed are # ordered by timestep, with shape `[num_timesteps, num_samples, # batch_shape, size, 1]` where `size` represents `latent_size` # or `observation_size` respectively. But timesteps are really # part of each probabilistic event, so we need to return a Tensor # of shape `[num_samples, batch_shape, num_timesteps, size]`. latents = tf . squeeze ( latents , - 1 ) latents = distribution_util . move_dimension ( latents , 0 , - 2 ) observations = tf . squeeze ( observations , - 1 ) observations = distribution_util . move_dimension ( observations , 0 , - 2 ) return latents , observations
Draw a joint sample from the prior over latents and observations.
Draw a joint sample from the prior over latents and observations .