code: def kalman_filter(
      transition_matrix,
      transition_mean,
      transition_cov,
      observation_matrix,
      observation_mean,
      observation_cov,
      initial_mean,
      initial_cov,
      y,
      mask,
      return_all=True,
  ):
      """Infers latent values using a parallel Kalman filter.
  
      This method computes filtered marginal means and covariances of a linear
      Gaussian state-space model using a parallel message-passing algorithm, as
      described by Sarkka and Garcia-Fernandez [1]. The inference process is
      formulated as a prefix-sum problem that can be efficiently computed by
      `tfp.math.scan_associative`, so that inference for a time series of length
      `num_timesteps` requires only `O(log(num_timesteps))` sequential steps.
  
      As with a naive sequential implementation, the total FLOP count scales
      linearly in `num_timesteps` (as `O(T + T/2 + T/4 + ...) = O(T)`), so this
      approach does not require extra resources in an asymptotic sense. However, it
      likely has a somewhat larger constant factor, so a sequential filter may be
      preferred when throughput rather than latency is the highest priority.
  
      Args:
        transition_matrix: float `Tensor` of shape
           `[num_timesteps, B1, .., BN, latent_size, latent_size]`.
        transition_mean: float `Tensor` of shape
           `[num_timesteps, B1, .., BN, latent_size]`.
        transition_cov: float `Tensor` of shape
           `[num_timesteps, B1, .., BN, latent_size, latent_size]`.
        observation_matrix: float `Tensor` of shape
           `[num_timesteps, B1, .., BN, observation_size, latent_size]`.
        observation_mean: float `Tensor` of shape
           `[num_timesteps, B1, .., BN, observation_size]`.
        observation_cov: float `Tensor` of shape
           `[num_timesteps, B1, .., BN, observation_size, observation_size]`.
        initial_mean: float `Tensor` of shape
           `[B1, .., BN, latent_size]`.
        initial_cov: float `Tensor` of shape
           `[B1, .., BN, latent_size, latent_size]`.
        y: float `Tensor` of shape
           `[num_timesteps, B1, .., BN, observation_size]`.
        mask: float `Tensor` of shape `[num_timesteps, B1, .., BN]`.
        return_all: Python `bool`, whether to compute log-likelihoods and
          predictive and observation distributions. If `False`, only
          `filtered_means` and `filtered_covs` are computed, and `None` is returned
          for the remaining values.
      Returns:
        log_likelihoods: float `Tensor` of shape `[num_timesteps, B1, .., BN]`, such
          that `log_likelihoods[t] = log p(y[t] | y[:t])`.
        filtered_means: float `Tensor` of shape
          `[num_timesteps, B1, .., BN, latent_size]`, such that
          `filtered_means[t] == E[x[t] | y[:t + 1]]`.
        filtered_covs: float `Tensor` of shape
          `[num_timesteps, B1, .., BN, latent_size, latent_size]`.
        predictive_means: float `Tensor` of shape
          `[num_timesteps, B1, .., BN, latent_size]`, such that
          `predictive_means[t] = E[x[t + 1] | y[:t + 1]]`.
        predictive_covs: float `Tensor` of shape
          `[num_timesteps, B1, .., BN, latent_size, latent_size]`.
        observation_means: float `Tensor` of shape
          `[num_timesteps, B1, .., BN, observation_size]`, such that
          `observation_means[t] = E[y[t] | y[:t]]`.
        observation_covs:float `Tensor` of shape
          `[num_timesteps, B1, .., BN, observation_size, observation_size]`.
  
      ### Mathematical Details
  
      The assumed model consists of latent state vectors
      `x[:num_timesteps, :latent_size]` and corresponding observed values
      `y[:num_timesteps, :observation_size]`, governed by the following dynamics:
  
      ```
      x[0] ~ MultivariateNormal(mean=initial_mean, cov=initial_cov)
      for t in range(num_timesteps - 1):
        x[t + 1] ~ MultivariateNormal(mean=matmul(transition_matrix[t],
                                                  x[t]) + transition_mean[t],
                                      cov=transition_cov[t])
      # Observed values `y[:num_timesteps]` defined at all timesteps.
      y ~ MultivariateNormal(mean=matmul(observation_matrix, x) + observation_mean,
                             cov=observation_cov)
      ```
  
      ### Tensor layout
  
      `Tensor` arguments are expected to have `num_timesteps` as their *leftmost*
      axis, preceding any batch dimensions. This layout is used
      for internal computations, so providing arguments in this form avoids the
      need for potentially-spurious transposition. The returned `Tensor`s also
      follow this layout, for the same reason. Note that this differs from the
      layout mandated by the `tfd.Distribution`
      API (and exposed by `tfd.LinearGaussianStateSpaceModel`), in which the time
      axis is to the right of any batch dimensions; it is the caller's
      responsibility to perform any needed transpositions.
  
      Tensor arguments may be specified with partial batch shape, i.e., with
      shape prefix `[num_timesteps, Bk, ..., BN]` for `k > 1`. They will be
      internally reshaped and broadcast to the full batch shape prefix
      `[num_timesteps, B1, ..., BN]`.
  
      ### References
  
      [1] Simo Sarkka and Angel F. Garcia-Fernandez. Temporal Parallelization of
          Bayesian Smoothers. _arXiv preprint arXiv:1905.13002_, 2019.
          https://arxiv.org/abs/1905.13002
  
      """
      with tf.name_scope("kalman_filter"):
          (
              time_indep,
              time_dep,
              observation,
          ) = broadcast_to_full_batch_shape(
              time_indep=TimeIndependentParameters(
                  initial_cov=initial_cov,
                  initial_scale_tril=None,
                  initial_mean=initial_mean,
              ),
              time_dep=TimeDependentParameters(
                  transition_matrix=transition_matrix,
                  transition_cov=transition_cov,
                  transition_scale_tril=None,
                  transition_mean=transition_mean,
                  observation_matrix=observation_matrix,
                  observation_cov=observation_cov,
                  observation_scale_tril=None,
                  observation_mean=observation_mean,
              ),
              observation=Observations(y, mask),
          )
  
          # Prevent any masked NaNs from leaking into gradients.
          if observation.mask is not None:
              observation = Observations(
                  y=tf.where(
                      observation.mask[..., None],
                      tf.zeros([], dtype=observation.y.dtype),
                      observation.y,
                  ),
                  mask=observation.mask,
              )
  
          # Run Kalman filter.
          filtered = scan_associative(
              combine_filter_elements,
              filter_elements(time_indep, time_dep, observation),
          )
          filtered_means = filtered.posterior_mean
          filtered_covs = filtered.posterior_cov
          log_likelihoods = None
          predicted_means, predicted_covs = None, None
          observation_means, observation_covs = None, None
          # Compute derived quantities (predictive distributions, likelihood, etc.).
          if return_all:
              predicted_means = _propagate_mean(
                  matrix=time_dep.transition_matrix,
                  mean=filtered_means,
                  added_mean=time_dep.transition_mean,
              )
              observation_means = _propagate_mean(
                  matrix=time_dep.observation_matrix,
                  mean=tf.concat(
                      [[time_indep.initial_mean], predicted_means[:-1]],
                      axis=0,
                  ),
                  added_mean=time_dep.observation_mean,
              )
              predicted_covs = _propagate_cov(
                  matrix=time_dep.transition_matrix,
                  cov=filtered_covs,
                  added_cov=time_dep.transition_cov,
              )
              observation_covs = _propagate_cov(
                  matrix=time_dep.observation_matrix,
                  cov=tf.concat(
                      [[time_indep.initial_cov], predicted_covs[:-1]],
                      axis=0,
                  ),
                  added_cov=time_dep.observation_cov,
              )
  
              log_likelihoods = _mvn_log_prob(
                  observation_means, observation_covs, observation.y
              )
              if observation.mask is not None:
                  log_likelihoods = tf.where(
                      observation.mask,
                      tf.zeros([], dtype=log_likelihoods.dtype),
                      log_likelihoods,
                  )
  
          return FilterResults(
              log_likelihoods,
              filtered_means,
              filtered_covs,
              predicted_means,
              predicted_covs,
              observation_means,
              observation_covs,
          )
  
imports: import collections
  import functools
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import mvn_tril
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.math import linalg
  from tensorflow_probability.python.math.scan_associative import scan_associative
label_desc: Use image node_id 10 for calling a global function with example usage: kalman_filter(transition_matrix, transition_mean, transition_cov, observation_matrix, observation_mean, observation_cov, initial_mean, initial_cov, y, mask, return_all) and returns: FilterResults