class_docstr: Implements Linear Dynamical System (LDS) as a distribution.
  
  The LDS is given by
  
  .. math::
      z_t = A_t l_{t-1} + b_t + \epsilon_t \\
      l_t = C_t l_{t-1} + g_t \nu
  
  where
  
  .. math::
      \epsilon_t = N(0, S_v) \\
      \nu = N(0, 1)
  
  :math:`A_t`, :math:`C_t` and :math:`g_t` are the emission, transition and
  innovation coefficients respectively. The residual terms are denoted
  by :math:`b_t`.
  
  The target :math:`z_t` can be :math:`d`-dimensional in which case
  
  .. math::
      A_t \in R^{d \times h}, b_t \in R^{d}, C_t \in R^{h \times h}, g_t \in R^{h}
  
  where :math:`h` is dimension of the latent state.
  
  Parameters
  ----------
  emission_coeff
      Tensor of shape (batch_size, seq_length, obs_dim, latent_dim)
  transition_coeff
      Tensor of shape (batch_size, seq_length, latent_dim, latent_dim)
  innovation_coeff
      Tensor of shape (batch_size, seq_length, latent_dim)
  noise_std
      Tensor of shape (batch_size, seq_length, obs_dim)
  residuals
      Tensor of shape (batch_size, seq_length, obs_dim)
  prior_mean
      Tensor of shape (batch_size, latent_dim)
  prior_cov
      Tensor of shape (batch_size, latent_dim, latent_dim)
  latent_dim
      Dimension of the latent state
  output_dim
      Dimension of the output
  seq_length
      Sequence length
  F
class_docstr_tok: ["Implements","Linear","Dynamical","System","(","LDS",")","as","a","distribution",".","The","LDS","is","given","by","..","math",":",":","z_t","=","A_t","l_","{","t-1","}","+","b_t","+","\\epsilon_t","\\\\","l_t","=","C_t","l_","{","t-1","}","+","g_t","\\nu","where","..","math",":",":","\\epsilon_t","=","N","(","0",",","S_v",")","\\\\","\\nu","=","N","(","0",",","1",")",":","math",":","`","A_t","`",",",":","math",":","`","C_t","`","and",":","math",":","`","g_t","`","are","the","emission",",","transition","and","innovation","coefficients","respectively",".","The","residual","terms","are","denoted","by",":","math",":","`","b_t","`",".","The","target",":","math",":","`","z_t","`","can","be",":","math",":","`","d","`","-dimensional","in","which","case","..","math",":",":","A_t","\\in","R^","{","d","\\times","h","}",",","b_t","\\in","R^","{","d","}",",","C_t","\\in","R^","{","h","\\times","h","}",",","g_t","\\in","R^","{","h","}","where",":","math",":","`","h","`","is","dimension","of","the","latent","state",".","Parameters","--","--","--","--","--","emission_coeff","Tensor","of","shape","(","batch_size",",","seq_length",",","obs_dim",",","latent_dim",")","transition_coeff","Tensor","of","shape","(","batch_size",",","seq_length",",","latent_dim",",","latent_dim",")","innovation_coeff","Tensor","of","shape","(","batch_size",",","seq_length",",","latent_dim",")","noise_std","Tensor","of","shape","(","batch_size",",","seq_length",",","obs_dim",")","residuals","Tensor","of","shape","(","batch_size",",","seq_length",",","obs_dim",")","prior_mean","Tensor","of","shape","(","batch_size",",","latent_dim",")","prior_cov","Tensor","of","shape","(","batch_size",",","latent_dim",",","latent_dim",")","latent_dim","Dimension","of","the","latent","state","output_dim","Dimension","of","the","output","seq_length","Sequence","length","F"]
docstr: Performs Kalman filtering given observations.
  
  Parameters
  ----------
  targets
      Observations, shape (batch_size, seq_length, output_dim)
  observed
      Flag tensor indicating which observations are genuine (1.0) and
      which are missing (0.0)
  
  Returns
  -------
  Tensor
      Log probabilities, shape (batch_size, seq_length)
  Tensor
      Mean of p(l_T | l_{T-1}), where T is seq_length, with shape
      (batch_size, latent_dim)
  Tensor
      Covariance of p(l_T | l_{T-1}), where T is seq_length, with shape
      (batch_size, latent_dim, latent_dim)
docstr_tok: ["Performs","Kalman","filtering","given","observations",".","Parameters","--","--","--","--","--","targets","Observations",",","shape","(","batch_size",",","seq_length",",","output_dim",")","observed","Flag","tensor","indicating","which","observations","are","genuine","(","1.0",")","and","which","are","missing","(","0.0",")","Returns","--","--","--","-","Tensor","Log","probabilities",",","shape","(","batch_size",",","seq_length",")","Tensor","Mean","of","p","(","l_T","|","l_","{","T-1","}",")",",","where","T","is","seq_length",",","with","shape","(","batch_size",",","latent_dim",")","Tensor","Covariance","of","p","(","l_T","|","l_","{","T-1","}",")",",","where","T","is","seq_length",",","with","shape","(","batch_size",",","latent_dim",",","latent_dim",")"]
code: def kalman_filter(
      self, targets: Tensor, observed: Tensor
  ) -> Tuple[Tensor, ...]:
      """
      Performs Kalman filtering given observations.
  
      Parameters
      ----------
      targets
          Observations, shape (batch_size, seq_length, output_dim)
      observed
          Flag tensor indicating which observations are genuine (1.0) and
          which are missing (0.0)
  
      Returns
      -------
      Tensor
          Log probabilities, shape (batch_size, seq_length)
      Tensor
          Mean of p(l_T | l_{T-1}), where T is seq_length, with shape
          (batch_size, latent_dim)
      Tensor
          Covariance of p(l_T | l_{T-1}), where T is seq_length, with shape
          (batch_size, latent_dim, latent_dim)
      """
      F = self.F
      # targets[t]: (batch_size, obs_dim)
      targets = _safe_split(
          targets,
          axis=1,
          num_outputs=self.seq_length,
          squeeze_axis=True,
      )
  
      log_p_seq = []
  
      mean = self.prior_mean
      cov = self.prior_cov
  
      observed = (
          _safe_split(
              observed,
              axis=1,
              num_outputs=self.seq_length,
              squeeze_axis=True,
          )
          if observed is not None
          else None
      )
  
      for t in range(self.seq_length):
          # Compute the filtered distribution
          #   p(l_t | z_1, ..., z_{t + 1})
          # and log - probability
          #   log p(z_t | z_0, z_{t - 1})
          filtered_mean, filtered_cov, log_p = kalman_filter_step(
              F,
              target=targets[t],
              prior_mean=mean,
              prior_cov=cov,
              emission_coeff=self.emission_coeff[t],
              residual=self.residuals[t],
              noise_std=self.noise_std[t],
              latent_dim=self.latent_dim,
              output_dim=self.output_dim,
          )
  
          log_p_seq.append(log_p.expand_dims(axis=1))
  
          # Mean of p(l_{t+1} | l_t)
          mean = F.linalg_gemm2(
              self.transition_coeff[t],
              (
                  filtered_mean.expand_dims(axis=-1)
                  if observed is None
                  else F.where(
                      observed[t], x=filtered_mean, y=mean
                  ).expand_dims(axis=-1)
              ),
          ).squeeze(axis=-1)
  
          # Covariance of p(l_{t+1} | l_t)
          cov = F.linalg_gemm2(
              self.transition_coeff[t],
              F.linalg_gemm2(
                  (
                      filtered_cov
                      if observed is None
                      else F.where(observed[t], x=filtered_cov, y=cov)
                  ),
                  self.transition_coeff[t],
                  transpose_b=True,
              ),
          ) + F.linalg_gemm2(
              self.innovation_coeff[t],
              self.innovation_coeff[t],
              transpose_a=True,
          )
  
      # Return sequence of log likelihoods, as well as
      # final mean and covariance of p(l_T | l_{T-1} where T is seq_length
      return F.concat(*log_p_seq, dim=1), mean, cov
  
imports: from typing import Optional, Tuple
  import mxnet
  import numpy
  from gluonts.core.component import validated
  from gluonts.mx import Tensor
  from gluonts.mx.linalg_util import jitter_cholesky
  from gluonts.mx.util import _broadcast_param, make_nd_diag
  from .None import Distribution, Gaussian, MultivariateGaussian
  from .distribution import getF
label_desc: Use image node_id 7 for calling the LDS obj's underlying member method code with example usage: obj.kalman_filter(targets, observed) and returns: F,  mean,  cov