code: def test_initial_state_broadcasts_over_batch(self):
      batch_shape = [4, 3]
      num_timesteps = 10
      num_features = 2
      drift_scale = 1.0
  
      initial_state_loc = self._build_placeholder([0.1, -0.2])
      initial_state_scale = self._build_placeholder([0.42, 0.314])
  
      design_matrix = self._build_placeholder(
          np.random.randn(
              *(batch_shape + [num_timesteps, num_features])
          )
      )
  
      initial_state_prior = mvn_diag.MultivariateNormalDiag(
          loc=initial_state_loc, scale_diag=initial_state_scale
      )
  
      ssm = DynamicLinearRegressionStateSpaceModel(
          num_timesteps=num_timesteps,
          design_matrix=design_matrix,
          drift_scale=drift_scale,
          initial_state_prior=initial_state_prior,
      )
  
      seed = test_util.test_seed(sampler_type="stateless")
      sample = ssm.sample(seed=seed)
  
      ll, means, _, _, _, _, _ = ssm.forward_filter(sample)
  
      self.assertAllEqual(
          batch_shape + [num_timesteps, 1], self.evaluate(sample).shape
      )
  
      self.assertAllEqual(
          batch_shape + [num_timesteps], self.evaluate(ll).shape
      )
  
      self.assertAllEqual(
          batch_shape + [num_timesteps, num_features],
          self.evaluate(means).shape,
      )
  
imports: import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.sts.components.dynamic_regression import DynamicLinearRegression
  from tensorflow_probability.python.sts.components.dynamic_regression import DynamicLinearRegressionStateSpaceModel
label_desc: Use image node_id 2 for calling the _DynamicLinearRegressionStateSpaceModelTest obj's underlying member method code with example usage: obj.test_initial_state_broadcasts_over_batch() without return types