code: def test_sample_stats(self):
      ndim = 3
      mdim = 2
      nsteps = 3
      nsamples = 10000
  
      seed = test_util.test_seed()
      _, seed = samplers.split_seed(seed, n=2, salt="")
      s = samplers.split_seed(seed, n=9, salt="")
      initial_mean = self.evaluate(
          _random_vector(ndim, (), dtype=self.dtype, seed=s[0])
      )
      initial_cov = self.evaluate(
          _random_variance(ndim, (), dtype=self.dtype, seed=s[1])
      )
      transition_matrix = self.evaluate(
          0.5
          * _random_matrix(
              ndim, ndim, (nsteps,), dtype=self.dtype, seed=s[2]
          )
      )
      transition_mean = self.evaluate(
          _random_vector(ndim, (nsteps,), dtype=self.dtype, seed=s[3])
      )
      transition_cov = self.evaluate(
          0.5
          * _random_variance(
              ndim, (nsteps,), dtype=self.dtype, seed=s[4]
          )
      )
      observation_matrix = self.evaluate(
          _random_matrix(
              mdim, ndim, (nsteps,), dtype=self.dtype, seed=s[5]
          )
      )
      observation_mean = self.evaluate(
          _random_vector(mdim, (nsteps,), dtype=self.dtype, seed=s[6])
      )
      observation_cov = self.evaluate(
          _random_variance(mdim, (nsteps,), dtype=self.dtype, seed=s[7])
      )
  
      x, y = self.evaluate(
          parallel_kalman_filter_lib.sample_walk(
              transition_matrix=transition_matrix,
              transition_scale_tril=tf.linalg.cholesky(transition_cov),
              observation_matrix=observation_matrix,
              observation_scale_tril=tf.linalg.cholesky(
                  observation_cov
              ),
              initial_mean=tf.broadcast_to(
                  initial_mean,
                  ps.concat(
                      [[nsamples], ps.shape(initial_mean)], axis=0
                  ),
              ),
              initial_scale_tril=tf.linalg.cholesky(initial_cov),
              transition_mean=transition_mean,
              observation_mean=observation_mean,
              seed=s[8],
          )
      )
      empirical_initial_mean = np.mean(x[0], axis=0)
      empirical_initial_cov = self.evaluate(
          sample_stats.covariance(x[0])
      )
      self.assertAllClose(
          empirical_initial_mean, initial_mean, atol=0.1
      )
      self.assertAllClose(empirical_initial_cov, initial_cov, atol=0.1)
  
      transition_residuals = self.evaluate(
          x[1:]
          - tf.linalg.matvec(
              transition_matrix[:-1, tf.newaxis, ...], x[:-1]
          )
      )
      empirical_transition_mean = np.mean(transition_residuals, axis=1)
      empirical_transition_cov = self.evaluate(
          sample_stats.covariance(transition_residuals, sample_axis=1)
      )
      # Checking the mean and cov of residuals implicitly also
      # checks the `transition_matrix` used to compute them.
      self.assertAllClose(
          empirical_transition_mean, transition_mean[:-1], atol=0.1
      )  # Typical error with 10k samples ~= 0.01-0.03
      self.assertAllClose(
          empirical_transition_cov, transition_cov[:-1], atol=0.1
      )
  
      observation_residuals = self.evaluate(
          y
          - tf.linalg.matvec(observation_matrix[:, tf.newaxis, ...], x)
      )
      empirical_observation_mean = np.mean(
          observation_residuals, axis=1
      )
      empirical_observation_cov = self.evaluate(
          sample_stats.covariance(observation_residuals, sample_axis=1)
      )
      self.assertAllClose(
          empirical_observation_mean, observation_mean, atol=0.1
      )
      self.assertAllClose(
          empirical_observation_cov, observation_cov, atol=0.1
      )
  
imports: import collections
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import linear_gaussian_ssm
  from tensorflow_probability.python.distributions import mvn_tril
  from tensorflow_probability.python.experimental.parallel_filter import parallel_kalman_filter_lib
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
  from tensorflow_probability.python.stats import sample_stats
label_desc: Use image node_id 5 for calling the _KalmanFilterTest obj's underlying member method code with example usage: obj.test_sample_stats() without return types