def build_pushforward_latents_step(get_observation_matrix_for_timestep,
                                   get_observation_noise_for_timestep):
  """Build a callable to push latent means/covs to observed means/covs.

  Args:
    get_observation_matrix_for_timestep: callable taking a timestep
      as an integer `Tensor` argument, and returning a `LinearOperator`
      of shape `[observation_size, observation_size]`.
    get_observation_noise_for_timestep: callable taking a timestep as
      an integer `Tensor` argument, and returning a
      `MultivariateNormalLinearOperator` of event shape
      `[observation_size]`.

  Returns:
    pushforward_latents_step: a callable that computes the observation mean and
    covariance at time `t`, given latent mean and covariance at time `t`.
  """

  def pushforward_latents_step(_, latent_t_mean_cov):
    """Loop body fn to pushforward latents to observations at a time step."""
    t, latent_mean, latent_cov = latent_t_mean_cov

    observation_matrix = get_observation_matrix_for_timestep(t)
    observation_noise = get_observation_noise_for_timestep(t)
    observation_mean = _propagate_mean(latent_mean,
                                       observation_matrix,
                                       observation_noise)
    observation_cov = _propagate_cov(latent_cov,
                                     observation_matrix,
                                     observation_noise)

    return (observation_mean, observation_cov)

  return pushforward_latents_step
def build_pushforward_latents_step(get_observation_matrix_for_timestep,
                                   get_observation_noise_for_timestep):
  """Build a callable to push latent means/covs to observed means/covs.

  Args:
    get_observation_matrix_for_timestep: callable taking a timestep
      as an integer `Tensor` argument, and returning a `LinearOperator`
      of shape `[observation_size, observation_size]`.
    get_observation_noise_for_timestep: callable taking a timestep as
      an integer `Tensor` argument, and returning a
      `MultivariateNormalLinearOperator` of event shape
      `[observation_size]`.

  Returns:
    pushforward_latents_step: a callable that computes the observation mean and
    covariance at time `t`, given latent mean and covariance at time `t`.
  """

  def pushforward_latents_step(_, latent_t_mean_cov):
    """Loop body fn to pushforward latents to observations at a time step."""
    t, latent_mean, latent_cov = latent_t_mean_cov

    observation_matrix = get_observation_matrix_for_timestep(t)
    observation_noise = get_observation_noise_for_timestep(t)
    observation_mean = _propagate_mean(latent_mean,
                                       observation_matrix,
                                       observation_noise)
    observation_cov = _propagate_cov(latent_cov,
                                     observation_matrix,
                                     observation_noise)

    return (observation_mean, observation_cov)

  return pushforward_latents_step
def build_pushforward_latents_step ( get_observation_matrix_for_timestep , get_observation_noise_for_timestep ) : def pushforward_latents_step ( _ , latent_t_mean_cov ) : """Loop body fn to pushforward latents to observations at a time step.""" t , latent_mean , latent_cov = latent_t_mean_cov observation_matrix = get_observation_matrix_for_timestep ( t ) observation_noise = get_observation_noise_for_timestep ( t ) observation_mean = _propagate_mean ( latent_mean , observation_matrix , observation_noise ) observation_cov = _propagate_cov ( latent_cov , observation_matrix , observation_noise ) return ( observation_mean , observation_cov ) return pushforward_latents_step
Build a callable to push latent means/covs to observed means/covs.

  Args:
    get_observation_matrix_for_timestep: callable taking a timestep
      as an integer `Tensor` argument, and returning a `LinearOperator`
      of shape `[observation_size, observation_size]`.
    get_observation_noise_for_timestep: callable taking a timestep as
      an integer `Tensor` argument, and returning a
      `MultivariateNormalLinearOperator` of event shape
      `[observation_size]`.

  Returns:
    pushforward_latents_step: a callable that computes the observation mean and
    covariance at time `t`, given latent mean and covariance at time `t`.
Build a callable to push latent means / covs to observed means / covs .