code: def test_broadcasting_correctness(self):
      seed = test_util.test_seed(sampler_type="stateless")
  
      # This test verifies that broadcasting of component parameters works as
      # expected. We construct a SSM with no batch shape, and test that when we
      # add it to another SSM of batch shape [3], we get the same model
      # as if we had explicitly broadcast the parameters of the first SSM before
      # adding.
  
      num_timesteps = 5
      transition_matrix = np.random.randn(2, 2)
      transition_noise_diag = np.exp(np.random.randn(2))
      observation_matrix = np.random.randn(1, 2)
      observation_noise_diag = np.exp(np.random.randn(1))
      initial_state_prior_diag = np.exp(np.random.randn(2))
  
      # First build the model in which we let AdditiveSSM do the broadcasting.
      batchless_ssm = linear_gaussian_ssm.LinearGaussianStateSpaceModel(
          num_timesteps=num_timesteps,
          transition_matrix=self._build_placeholder(transition_matrix),
          transition_noise=mvn_diag.MultivariateNormalDiag(
              scale_diag=self._build_placeholder(transition_noise_diag)
          ),
          observation_matrix=self._build_placeholder(
              observation_matrix
          ),
          observation_noise=mvn_diag.MultivariateNormalDiag(
              scale_diag=self._build_placeholder(observation_noise_diag)
          ),
          initial_state_prior=mvn_diag.MultivariateNormalDiag(
              scale_diag=self._build_placeholder(
                  initial_state_prior_diag
              )
          ),
      )
      another_ssm = self._dummy_model(
          num_timesteps=num_timesteps, latent_size=4, batch_shape=[3]
      )
      broadcast_additive_ssm = AdditiveStateSpaceModel(
          [batchless_ssm, another_ssm]
      )
  
      # Next try doing our own broadcasting explicitly.
      broadcast_vector = np.ones([3, 1])
      broadcast_matrix = np.ones([3, 1, 1])
      batch_ssm = linear_gaussian_ssm.LinearGaussianStateSpaceModel(
          num_timesteps=num_timesteps,
          transition_matrix=self._build_placeholder(
              transition_matrix * broadcast_matrix
          ),
          transition_noise=mvn_diag.MultivariateNormalDiag(
              scale_diag=self._build_placeholder(
                  transition_noise_diag * broadcast_vector
              )
          ),
          observation_matrix=self._build_placeholder(
              observation_matrix * broadcast_matrix
          ),
          observation_noise=mvn_diag.MultivariateNormalDiag(
              scale_diag=self._build_placeholder(
                  observation_noise_diag * broadcast_vector
              )
          ),
          initial_state_prior=mvn_diag.MultivariateNormalDiag(
              scale_diag=self._build_placeholder(
                  initial_state_prior_diag * broadcast_vector
              )
          ),
      )
      manual_additive_ssm = AdditiveStateSpaceModel(
          [batch_ssm, another_ssm]
      )
  
      # Both additive SSMs define the same model, so they should give the same
      # log_probs.
      y = self.evaluate(broadcast_additive_ssm.sample(seed=seed))
      self.assertAllClose(
          self.evaluate(broadcast_additive_ssm.log_prob(y)),
          self.evaluate(manual_additive_ssm.log_prob(y)),
      )
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import linear_gaussian_ssm
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.sts.components.local_linear_trend import LocalLinearTrendStateSpaceModel
  from tensorflow_probability.python.sts.components.sum import AdditiveStateSpaceModel
label_desc: Use image node_id 8 for calling the _AdditiveStateSpaceModelTest obj's underlying member method code with example usage: obj.test_broadcasting_correctness() without return types