tensorflow_probability/python/distributions/linear_gaussian_ssm.py
def forward_filter(self, x, mask=None):
    """Run a Kalman filter over a provided sequence of outputs.

    Note that the returned values `filtered_means`, `predicted_means`, and
    `observation_means` depend on the observed time series `x`, while the
    corresponding covariances are independent of the observed series; i.e., they
    depend only on the model itself. This means that the mean values have shape
    `concat([sample_shape(x), batch_shape, [num_timesteps,
    {latent/observation}_size]])`, while the covariances have shape
    `concat[(batch_shape, [num_timesteps, {latent/observation}_size,
    {latent/observation}_size]])`, which does not depend on the sample shape.

    Args:
      x: a float-type `Tensor` with rightmost dimensions
        `[num_timesteps, observation_size]` matching
        `self.event_shape`. Additional dimensions must match or be
        broadcastable to `self.batch_shape`; any further dimensions
        are interpreted as a sample shape.
      mask: optional bool-type `Tensor` with rightmost dimension
        `[num_timesteps]`; `True` values specify that the value of `x`
        at that timestep is masked, i.e., not conditioned on. Additional
        dimensions must match or be broadcastable to `self.batch_shape`; any
        further dimensions must match or be broadcastable to the sample
        shape of `x`.
        Default value: `None`.

    Returns:
      log_likelihoods: Per-timestep log marginal likelihoods `log
        p(x_t | x_{:t-1})` evaluated at the input `x`, as a `Tensor`
        of shape `sample_shape(x) + batch_shape + [num_timesteps].`
      filtered_means: Means of the per-timestep filtered marginal
         distributions p(z_t | x_{:t}), as a Tensor of shape
        `sample_shape(x) + batch_shape + [num_timesteps, latent_size]`.
      filtered_covs: Covariances of the per-timestep filtered marginal
         distributions p(z_t | x_{:t}), as a Tensor of shape
        `sample_shape(mask) + batch_shape + [num_timesteps, latent_size,
        latent_size]`. Note that the covariances depend only on the model and
        the mask, not on the data, so this may have fewer dimensions than
        `filtered_means`.
      predicted_means: Means of the per-timestep predictive
         distributions over latent states, p(z_{t+1} | x_{:t}), as a
         Tensor of shape `sample_shape(x) + batch_shape +
         [num_timesteps, latent_size]`.
      predicted_covs: Covariances of the per-timestep predictive
         distributions over latent states, p(z_{t+1} | x_{:t}), as a
         Tensor of shape `sample_shape(mask) + batch_shape +
         [num_timesteps, latent_size, latent_size]`. Note that the covariances
         depend only on the model and the mask, not on the data, so this may
         have fewer dimensions than `predicted_means`.
      observation_means: Means of the per-timestep predictive
         distributions over observations, p(x_{t} | x_{:t-1}), as a
         Tensor of shape `sample_shape(x) + batch_shape +
         [num_timesteps, observation_size]`.
      observation_covs: Covariances of the per-timestep predictive
         distributions over observations, p(x_{t} | x_{:t-1}), as a
         Tensor of shape `sample_shape(mask) + batch_shape + [num_timesteps,
         observation_size, observation_size]`. Note that the covariances depend
         only on the model and the mask, not on the data, so this may have fewer
         dimensions than `observation_means`.
    """

    with tf.name_scope("forward_filter"):
      x = tf.convert_to_tensor(value=x, name="x")
      if mask is not None:
        mask = tf.convert_to_tensor(value=mask, name="mask", dtype_hint=tf.bool)

      # Check event shape statically if possible
      check_x_shape_op = _check_equal_shape(
          "x", x.shape[-2:], tf.shape(input=x)[-2:],
          self.event_shape, self.event_shape_tensor())
      check_mask_dims_op = None
      check_mask_shape_op = None
      if mask is not None:
        if (tensorshape_util.rank(mask.shape) is None or
            tensorshape_util.rank(x.shape) is None):
          check_mask_dims_op = assert_util.assert_greater_equal(
              tf.rank(x),
              tf.rank(mask),
              message=("mask cannot have higher rank than x!"))
        elif tensorshape_util.rank(mask.shape) > tensorshape_util.rank(x.shape):
          raise ValueError(
              "mask cannot have higher rank than x! ({} vs {})".format(
                  tensorshape_util.rank(mask.shape),
                  tensorshape_util.rank(x.shape)))
        check_mask_shape_op = _check_equal_shape(
            "mask", mask.shape[-1:], tf.shape(input=mask)[-1:],
            self.event_shape[-2:-1], self.event_shape_tensor()[-2:-1])
      if self.validate_args:
        runtime_assertions = self.runtime_assertions
        if check_x_shape_op is not None:
          runtime_assertions += [check_x_shape_op]
        if check_mask_shape_op is not None:
          runtime_assertions += [check_mask_shape_op]
        if check_mask_dims_op is not None:
          runtime_assertions += [check_mask_dims_op]
        with tf.control_dependencies(runtime_assertions):
          x = tf.identity(x)

      # Get the full output sample_shape + batch shape. Usually
      # this will just be x[:-2], i.e. the input shape excluding
      # event shape. But users can specify inputs that broadcast
      # batch dimensions, so we need to broadcast this against
      # self.batch_shape.
      if tensorshape_util.is_fully_defined(
          self.batch_shape) and tensorshape_util.is_fully_defined(x.shape):
        sample_and_batch_shape = tf.broadcast_static_shape(
            x.shape[:-2], self.batch_shape)
      else:
        sample_and_batch_shape = tf.broadcast_dynamic_shape(
            tf.shape(input=x)[:-2], self.batch_shape_tensor())

      # Get the full output shape for covariances. The posterior variances
      # in a LGSSM depend only on the model params (batch shape) and on the
      # missingness pattern (mask shape), so in general this may be smaller
      # than the full `sample_and_batch_shape`.
      if mask is None:
        mask_sample_and_batch_shape = self.batch_shape_tensor()
      else:
        if (tensorshape_util.is_fully_defined(self.batch_shape) and
            tensorshape_util.is_fully_defined(mask.shape)):
          mask_sample_and_batch_shape = tf.broadcast_static_shape(
              mask.shape[:-1], self.batch_shape)
        else:
          mask_sample_and_batch_shape = tf.broadcast_dynamic_shape(
              tf.shape(input=mask)[:-1], self.batch_shape_tensor())

      # To scan over timesteps we need to move `num_timsteps` from the
      # event shape to the initial dimension of the tensor.
      x = distribution_util.move_dimension(x, -2, 0)
      if mask is not None:
        mask = distribution_util.move_dimension(mask, -1, 0)

      # Observations are assumed to be vectors, but we add a dummy
      # extra dimension to allow us to use `matmul` throughout.
      x = x[..., tf.newaxis]
      if mask is not None:
        # Align mask.shape with x.shape, including a unit dimension to broadcast
        # against `observation_size`.
        mask = mask[..., tf.newaxis, tf.newaxis]

      # Initialize filtering distribution from the prior. The mean in
      # a Kalman filter depends on data, so should match the full
      # sample and batch shape. The covariance is data-independent, so
      # only has batch shape.
      prior_mean = _broadcast_to_shape(
          self.initial_state_prior.mean()[..., tf.newaxis],
          tf.concat([sample_and_batch_shape,
                     [self.latent_size, 1]], axis=0))
      prior_cov = _broadcast_to_shape(
          self.initial_state_prior.covariance(),
          tf.concat([mask_sample_and_batch_shape,
                     [self.latent_size, self.latent_size]], axis=0))

      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_mean = _propagate_mean(prior_mean,
                                                 initial_observation_matrix,
                                                 initial_observation_noise)
      initial_observation_cov = _propagate_cov(prior_cov,
                                               initial_observation_matrix,
                                               initial_observation_noise)

      initial_state = KalmanFilterState(
          predicted_mean=prior_mean,
          predicted_cov=prior_cov,
          filtered_mean=prior_mean,  # establishes shape, value ignored
          filtered_cov=prior_cov,  # establishes shape, value ignored
          observation_mean=initial_observation_mean,
          observation_cov=initial_observation_cov,
          log_marginal_likelihood=tf.zeros(
              shape=sample_and_batch_shape, dtype=self.dtype),
          timestep=tf.convert_to_tensor(
              value=self.initial_step, dtype=tf.int32, name="initial_step"))

      update_step_fn = build_kalman_filter_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)

      filter_states = tf.scan(update_step_fn,
                              elems=x if mask is None else (x, mask),
                              initializer=initial_state)

      log_likelihoods = distribution_util.move_dimension(
          filter_states.log_marginal_likelihood, 0, -1)

      # Move the time dimension back into the event shape.
      filtered_means = distribution_util.move_dimension(
          filter_states.filtered_mean[..., 0], 0, -2)
      filtered_covs = distribution_util.move_dimension(
          filter_states.filtered_cov, 0, -3)
      predicted_means = distribution_util.move_dimension(
          filter_states.predicted_mean[..., 0], 0, -2)
      predicted_covs = distribution_util.move_dimension(
          filter_states.predicted_cov, 0, -3)
      observation_means = distribution_util.move_dimension(
          filter_states.observation_mean[..., 0], 0, -2)
      observation_covs = distribution_util.move_dimension(
          filter_states.observation_cov, 0, -3)
      # We could directly construct the batch Distributions
      # filtered_marginals = tfd.MultivariateNormalFullCovariance(
      #      filtered_means, filtered_covs)
      # predicted_marginals = tfd.MultivariateNormalFullCovariance(
      #      predicted_means, predicted_covs)
      # but we choose not to: returning the raw means and covariances
      # saves computation in Eager mode (avoiding an immediate
      # Cholesky factorization that the user may not want) and aids
      # debugging of numerical issues.

      return (log_likelihoods,
              filtered_means, filtered_covs,
              predicted_means, predicted_covs,
              observation_means, observation_covs)
def forward_filter(self, x, mask=None):
    """Run a Kalman filter over a provided sequence of outputs.

    Note that the returned values `filtered_means`, `predicted_means`, and
    `observation_means` depend on the observed time series `x`, while the
    corresponding covariances are independent of the observed series; i.e., they
    depend only on the model itself. This means that the mean values have shape
    `concat([sample_shape(x), batch_shape, [num_timesteps,
    {latent/observation}_size]])`, while the covariances have shape
    `concat[(batch_shape, [num_timesteps, {latent/observation}_size,
    {latent/observation}_size]])`, which does not depend on the sample shape.

    Args:
      x: a float-type `Tensor` with rightmost dimensions
        `[num_timesteps, observation_size]` matching
        `self.event_shape`. Additional dimensions must match or be
        broadcastable to `self.batch_shape`; any further dimensions
        are interpreted as a sample shape.
      mask: optional bool-type `Tensor` with rightmost dimension
        `[num_timesteps]`; `True` values specify that the value of `x`
        at that timestep is masked, i.e., not conditioned on. Additional
        dimensions must match or be broadcastable to `self.batch_shape`; any
        further dimensions must match or be broadcastable to the sample
        shape of `x`.
        Default value: `None`.

    Returns:
      log_likelihoods: Per-timestep log marginal likelihoods `log
        p(x_t | x_{:t-1})` evaluated at the input `x`, as a `Tensor`
        of shape `sample_shape(x) + batch_shape + [num_timesteps].`
      filtered_means: Means of the per-timestep filtered marginal
         distributions p(z_t | x_{:t}), as a Tensor of shape
        `sample_shape(x) + batch_shape + [num_timesteps, latent_size]`.
      filtered_covs: Covariances of the per-timestep filtered marginal
         distributions p(z_t | x_{:t}), as a Tensor of shape
        `sample_shape(mask) + batch_shape + [num_timesteps, latent_size,
        latent_size]`. Note that the covariances depend only on the model and
        the mask, not on the data, so this may have fewer dimensions than
        `filtered_means`.
      predicted_means: Means of the per-timestep predictive
         distributions over latent states, p(z_{t+1} | x_{:t}), as a
         Tensor of shape `sample_shape(x) + batch_shape +
         [num_timesteps, latent_size]`.
      predicted_covs: Covariances of the per-timestep predictive
         distributions over latent states, p(z_{t+1} | x_{:t}), as a
         Tensor of shape `sample_shape(mask) + batch_shape +
         [num_timesteps, latent_size, latent_size]`. Note that the covariances
         depend only on the model and the mask, not on the data, so this may
         have fewer dimensions than `predicted_means`.
      observation_means: Means of the per-timestep predictive
         distributions over observations, p(x_{t} | x_{:t-1}), as a
         Tensor of shape `sample_shape(x) + batch_shape +
         [num_timesteps, observation_size]`.
      observation_covs: Covariances of the per-timestep predictive
         distributions over observations, p(x_{t} | x_{:t-1}), as a
         Tensor of shape `sample_shape(mask) + batch_shape + [num_timesteps,
         observation_size, observation_size]`. Note that the covariances depend
         only on the model and the mask, not on the data, so this may have fewer
         dimensions than `observation_means`.
    """

    with tf.name_scope("forward_filter"):
      x = tf.convert_to_tensor(value=x, name="x")
      if mask is not None:
        mask = tf.convert_to_tensor(value=mask, name="mask", dtype_hint=tf.bool)

      # Check event shape statically if possible
      check_x_shape_op = _check_equal_shape(
          "x", x.shape[-2:], tf.shape(input=x)[-2:],
          self.event_shape, self.event_shape_tensor())
      check_mask_dims_op = None
      check_mask_shape_op = None
      if mask is not None:
        if (tensorshape_util.rank(mask.shape) is None or
            tensorshape_util.rank(x.shape) is None):
          check_mask_dims_op = assert_util.assert_greater_equal(
              tf.rank(x),
              tf.rank(mask),
              message=("mask cannot have higher rank than x!"))
        elif tensorshape_util.rank(mask.shape) > tensorshape_util.rank(x.shape):
          raise ValueError(
              "mask cannot have higher rank than x! ({} vs {})".format(
                  tensorshape_util.rank(mask.shape),
                  tensorshape_util.rank(x.shape)))
        check_mask_shape_op = _check_equal_shape(
            "mask", mask.shape[-1:], tf.shape(input=mask)[-1:],
            self.event_shape[-2:-1], self.event_shape_tensor()[-2:-1])
      if self.validate_args:
        runtime_assertions = self.runtime_assertions
        if check_x_shape_op is not None:
          runtime_assertions += [check_x_shape_op]
        if check_mask_shape_op is not None:
          runtime_assertions += [check_mask_shape_op]
        if check_mask_dims_op is not None:
          runtime_assertions += [check_mask_dims_op]
        with tf.control_dependencies(runtime_assertions):
          x = tf.identity(x)

      # Get the full output sample_shape + batch shape. Usually
      # this will just be x[:-2], i.e. the input shape excluding
      # event shape. But users can specify inputs that broadcast
      # batch dimensions, so we need to broadcast this against
      # self.batch_shape.
      if tensorshape_util.is_fully_defined(
          self.batch_shape) and tensorshape_util.is_fully_defined(x.shape):
        sample_and_batch_shape = tf.broadcast_static_shape(
            x.shape[:-2], self.batch_shape)
      else:
        sample_and_batch_shape = tf.broadcast_dynamic_shape(
            tf.shape(input=x)[:-2], self.batch_shape_tensor())

      # Get the full output shape for covariances. The posterior variances
      # in a LGSSM depend only on the model params (batch shape) and on the
      # missingness pattern (mask shape), so in general this may be smaller
      # than the full `sample_and_batch_shape`.
      if mask is None:
        mask_sample_and_batch_shape = self.batch_shape_tensor()
      else:
        if (tensorshape_util.is_fully_defined(self.batch_shape) and
            tensorshape_util.is_fully_defined(mask.shape)):
          mask_sample_and_batch_shape = tf.broadcast_static_shape(
              mask.shape[:-1], self.batch_shape)
        else:
          mask_sample_and_batch_shape = tf.broadcast_dynamic_shape(
              tf.shape(input=mask)[:-1], self.batch_shape_tensor())

      # To scan over timesteps we need to move `num_timsteps` from the
      # event shape to the initial dimension of the tensor.
      x = distribution_util.move_dimension(x, -2, 0)
      if mask is not None:
        mask = distribution_util.move_dimension(mask, -1, 0)

      # Observations are assumed to be vectors, but we add a dummy
      # extra dimension to allow us to use `matmul` throughout.
      x = x[..., tf.newaxis]
      if mask is not None:
        # Align mask.shape with x.shape, including a unit dimension to broadcast
        # against `observation_size`.
        mask = mask[..., tf.newaxis, tf.newaxis]

      # Initialize filtering distribution from the prior. The mean in
      # a Kalman filter depends on data, so should match the full
      # sample and batch shape. The covariance is data-independent, so
      # only has batch shape.
      prior_mean = _broadcast_to_shape(
          self.initial_state_prior.mean()[..., tf.newaxis],
          tf.concat([sample_and_batch_shape,
                     [self.latent_size, 1]], axis=0))
      prior_cov = _broadcast_to_shape(
          self.initial_state_prior.covariance(),
          tf.concat([mask_sample_and_batch_shape,
                     [self.latent_size, self.latent_size]], axis=0))

      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_mean = _propagate_mean(prior_mean,
                                                 initial_observation_matrix,
                                                 initial_observation_noise)
      initial_observation_cov = _propagate_cov(prior_cov,
                                               initial_observation_matrix,
                                               initial_observation_noise)

      initial_state = KalmanFilterState(
          predicted_mean=prior_mean,
          predicted_cov=prior_cov,
          filtered_mean=prior_mean,  # establishes shape, value ignored
          filtered_cov=prior_cov,  # establishes shape, value ignored
          observation_mean=initial_observation_mean,
          observation_cov=initial_observation_cov,
          log_marginal_likelihood=tf.zeros(
              shape=sample_and_batch_shape, dtype=self.dtype),
          timestep=tf.convert_to_tensor(
              value=self.initial_step, dtype=tf.int32, name="initial_step"))

      update_step_fn = build_kalman_filter_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)

      filter_states = tf.scan(update_step_fn,
                              elems=x if mask is None else (x, mask),
                              initializer=initial_state)

      log_likelihoods = distribution_util.move_dimension(
          filter_states.log_marginal_likelihood, 0, -1)

      # Move the time dimension back into the event shape.
      filtered_means = distribution_util.move_dimension(
          filter_states.filtered_mean[..., 0], 0, -2)
      filtered_covs = distribution_util.move_dimension(
          filter_states.filtered_cov, 0, -3)
      predicted_means = distribution_util.move_dimension(
          filter_states.predicted_mean[..., 0], 0, -2)
      predicted_covs = distribution_util.move_dimension(
          filter_states.predicted_cov, 0, -3)
      observation_means = distribution_util.move_dimension(
          filter_states.observation_mean[..., 0], 0, -2)
      observation_covs = distribution_util.move_dimension(
          filter_states.observation_cov, 0, -3)
      # We could directly construct the batch Distributions
      # filtered_marginals = tfd.MultivariateNormalFullCovariance(
      #      filtered_means, filtered_covs)
      # predicted_marginals = tfd.MultivariateNormalFullCovariance(
      #      predicted_means, predicted_covs)
      # but we choose not to: returning the raw means and covariances
      # saves computation in Eager mode (avoiding an immediate
      # Cholesky factorization that the user may not want) and aids
      # debugging of numerical issues.

      return (log_likelihoods,
              filtered_means, filtered_covs,
              predicted_means, predicted_covs,
              observation_means, observation_covs)
def forward_filter ( self , x , mask = None ) : with tf . name_scope ( "forward_filter" ) : x = tf . convert_to_tensor ( value = x , name = "x" ) if mask is not None : mask = tf . convert_to_tensor ( value = mask , name = "mask" , dtype_hint = tf . bool ) # Check event shape statically if possible check_x_shape_op = _check_equal_shape ( "x" , x . shape [ - 2 : ] , tf . shape ( input = x ) [ - 2 : ] , self . event_shape , self . event_shape_tensor ( ) ) check_mask_dims_op = None check_mask_shape_op = None if mask is not None : if ( tensorshape_util . rank ( mask . shape ) is None or tensorshape_util . rank ( x . shape ) is None ) : check_mask_dims_op = assert_util . assert_greater_equal ( tf . rank ( x ) , tf . rank ( mask ) , message = ( "mask cannot have higher rank than x!" ) ) elif tensorshape_util . rank ( mask . shape ) > tensorshape_util . rank ( x . shape ) : raise ValueError ( "mask cannot have higher rank than x! ({} vs {})" . format ( tensorshape_util . rank ( mask . shape ) , tensorshape_util . rank ( x . shape ) ) ) check_mask_shape_op = _check_equal_shape ( "mask" , mask . shape [ - 1 : ] , tf . shape ( input = mask ) [ - 1 : ] , self . event_shape [ - 2 : - 1 ] , self . event_shape_tensor ( ) [ - 2 : - 1 ] ) if self . validate_args : runtime_assertions = self . runtime_assertions if check_x_shape_op is not None : runtime_assertions += [ check_x_shape_op ] if check_mask_shape_op is not None : runtime_assertions += [ check_mask_shape_op ] if check_mask_dims_op is not None : runtime_assertions += [ check_mask_dims_op ] with tf . control_dependencies ( runtime_assertions ) : x = tf . identity ( x ) # Get the full output sample_shape + batch shape. Usually # this will just be x[:-2], i.e. the input shape excluding # event shape. But users can specify inputs that broadcast # batch dimensions, so we need to broadcast this against # self.batch_shape. if tensorshape_util . is_fully_defined ( self . batch_shape ) and tensorshape_util . is_fully_defined ( x . shape ) : sample_and_batch_shape = tf . broadcast_static_shape ( x . shape [ : - 2 ] , self . batch_shape ) else : sample_and_batch_shape = tf . broadcast_dynamic_shape ( tf . shape ( input = x ) [ : - 2 ] , self . batch_shape_tensor ( ) ) # Get the full output shape for covariances. The posterior variances # in a LGSSM depend only on the model params (batch shape) and on the # missingness pattern (mask shape), so in general this may be smaller # than the full `sample_and_batch_shape`. if mask is None : mask_sample_and_batch_shape = self . batch_shape_tensor ( ) else : if ( tensorshape_util . is_fully_defined ( self . batch_shape ) and tensorshape_util . is_fully_defined ( mask . shape ) ) : mask_sample_and_batch_shape = tf . broadcast_static_shape ( mask . shape [ : - 1 ] , self . batch_shape ) else : mask_sample_and_batch_shape = tf . broadcast_dynamic_shape ( tf . shape ( input = mask ) [ : - 1 ] , self . batch_shape_tensor ( ) ) # To scan over timesteps we need to move `num_timsteps` from the # event shape to the initial dimension of the tensor. x = distribution_util . move_dimension ( x , - 2 , 0 ) if mask is not None : mask = distribution_util . move_dimension ( mask , - 1 , 0 ) # Observations are assumed to be vectors, but we add a dummy # extra dimension to allow us to use `matmul` throughout. x = x [ ... , tf . newaxis ] if mask is not None : # Align mask.shape with x.shape, including a unit dimension to broadcast # against `observation_size`. mask = mask [ ... , tf . newaxis , tf . newaxis ] # Initialize filtering distribution from the prior. The mean in # a Kalman filter depends on data, so should match the full # sample and batch shape. The covariance is data-independent, so # only has batch shape. prior_mean = _broadcast_to_shape ( self . initial_state_prior . mean ( ) [ ... , tf . newaxis ] , tf . concat ( [ sample_and_batch_shape , [ self . latent_size , 1 ] ] , axis = 0 ) ) prior_cov = _broadcast_to_shape ( self . initial_state_prior . covariance ( ) , tf . concat ( [ mask_sample_and_batch_shape , [ self . latent_size , self . latent_size ] ] , axis = 0 ) ) 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_mean = _propagate_mean ( prior_mean , initial_observation_matrix , initial_observation_noise ) initial_observation_cov = _propagate_cov ( prior_cov , initial_observation_matrix , initial_observation_noise ) initial_state = KalmanFilterState ( predicted_mean = prior_mean , predicted_cov = prior_cov , filtered_mean = prior_mean , # establishes shape, value ignored filtered_cov = prior_cov , # establishes shape, value ignored observation_mean = initial_observation_mean , observation_cov = initial_observation_cov , log_marginal_likelihood = tf . zeros ( shape = sample_and_batch_shape , dtype = self . dtype ) , timestep = tf . convert_to_tensor ( value = self . initial_step , dtype = tf . int32 , name = "initial_step" ) ) update_step_fn = build_kalman_filter_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 ) filter_states = tf . scan ( update_step_fn , elems = x if mask is None else ( x , mask ) , initializer = initial_state ) log_likelihoods = distribution_util . move_dimension ( filter_states . log_marginal_likelihood , 0 , - 1 ) # Move the time dimension back into the event shape. filtered_means = distribution_util . move_dimension ( filter_states . filtered_mean [ ... , 0 ] , 0 , - 2 ) filtered_covs = distribution_util . move_dimension ( filter_states . filtered_cov , 0 , - 3 ) predicted_means = distribution_util . move_dimension ( filter_states . predicted_mean [ ... , 0 ] , 0 , - 2 ) predicted_covs = distribution_util . move_dimension ( filter_states . predicted_cov , 0 , - 3 ) observation_means = distribution_util . move_dimension ( filter_states . observation_mean [ ... , 0 ] , 0 , - 2 ) observation_covs = distribution_util . move_dimension ( filter_states . observation_cov , 0 , - 3 ) # We could directly construct the batch Distributions # filtered_marginals = tfd.MultivariateNormalFullCovariance( #      filtered_means, filtered_covs) # predicted_marginals = tfd.MultivariateNormalFullCovariance( #      predicted_means, predicted_covs) # but we choose not to: returning the raw means and covariances # saves computation in Eager mode (avoiding an immediate # Cholesky factorization that the user may not want) and aids # debugging of numerical issues. return ( log_likelihoods , filtered_means , filtered_covs , predicted_means , predicted_covs , observation_means , observation_covs )
Run a Kalman filter over a provided sequence of outputs.

    Note that the returned values `filtered_means`, `predicted_means`, and
    `observation_means` depend on the observed time series `x`, while the
    corresponding covariances are independent of the observed series; i.e., they
    depend only on the model itself. This means that the mean values have shape
    `concat([sample_shape(x), batch_shape, [num_timesteps,
    {latent/observation}_size]])`, while the covariances have shape
    `concat[(batch_shape, [num_timesteps, {latent/observation}_size,
    {latent/observation}_size]])`, which does not depend on the sample shape.

    Args:
      x: a float-type `Tensor` with rightmost dimensions
        `[num_timesteps, observation_size]` matching
        `self.event_shape`. Additional dimensions must match or be
        broadcastable to `self.batch_shape`; any further dimensions
        are interpreted as a sample shape.
      mask: optional bool-type `Tensor` with rightmost dimension
        `[num_timesteps]`; `True` values specify that the value of `x`
        at that timestep is masked, i.e., not conditioned on. Additional
        dimensions must match or be broadcastable to `self.batch_shape`; any
        further dimensions must match or be broadcastable to the sample
        shape of `x`.
        Default value: `None`.

    Returns:
      log_likelihoods: Per-timestep log marginal likelihoods `log
        p(x_t | x_{:t-1})` evaluated at the input `x`, as a `Tensor`
        of shape `sample_shape(x) + batch_shape + [num_timesteps].`
      filtered_means: Means of the per-timestep filtered marginal
         distributions p(z_t | x_{:t}), as a Tensor of shape
        `sample_shape(x) + batch_shape + [num_timesteps, latent_size]`.
      filtered_covs: Covariances of the per-timestep filtered marginal
         distributions p(z_t | x_{:t}), as a Tensor of shape
        `sample_shape(mask) + batch_shape + [num_timesteps, latent_size,
        latent_size]`. Note that the covariances depend only on the model and
        the mask, not on the data, so this may have fewer dimensions than
        `filtered_means`.
      predicted_means: Means of the per-timestep predictive
         distributions over latent states, p(z_{t+1} | x_{:t}), as a
         Tensor of shape `sample_shape(x) + batch_shape +
         [num_timesteps, latent_size]`.
      predicted_covs: Covariances of the per-timestep predictive
         distributions over latent states, p(z_{t+1} | x_{:t}), as a
         Tensor of shape `sample_shape(mask) + batch_shape +
         [num_timesteps, latent_size, latent_size]`. Note that the covariances
         depend only on the model and the mask, not on the data, so this may
         have fewer dimensions than `predicted_means`.
      observation_means: Means of the per-timestep predictive
         distributions over observations, p(x_{t} | x_{:t-1}), as a
         Tensor of shape `sample_shape(x) + batch_shape +
         [num_timesteps, observation_size]`.
      observation_covs: Covariances of the per-timestep predictive
         distributions over observations, p(x_{t} | x_{:t-1}), as a
         Tensor of shape `sample_shape(mask) + batch_shape + [num_timesteps,
         observation_size, observation_size]`. Note that the covariances depend
         only on the model and the mask, not on the data, so this may have fewer
         dimensions than `observation_means`.
Run a Kalman filter over a provided sequence of outputs .