code: def test_batch_shape(self):
      batch_shape = [3, 2]
      partial_batch_shape = [2]
      seed = test_util.test_seed(sampler_type="stateless")
  
      num_seasons = 24
      initial_state_prior = mvn_diag.MultivariateNormalDiag(
          scale_diag=self._build_placeholder(
              np.exp(
                  np.random.randn(
                      *(partial_batch_shape + [num_seasons - 1])
                  )
              )
          )
      )
      drift_scale = self._build_placeholder(
          np.exp(np.random.randn(*batch_shape))
      )
      observation_noise_scale = self._build_placeholder(
          np.exp(np.random.randn(*partial_batch_shape))
      )
  
      ssm = ConstrainedSeasonalStateSpaceModel(
          num_timesteps=9,
          num_seasons=24,
          num_steps_per_season=2,
          drift_scale=drift_scale,
          observation_noise_scale=observation_noise_scale,
          initial_state_prior=initial_state_prior,
      )
  
      # First check that the model's batch shape is the broadcast batch shape
      # of parameters, as expected.
  
      self.assertAllEqual(
          self.evaluate(ssm.batch_shape_tensor()), batch_shape
      )
      y_ = self.evaluate(ssm.sample(seed=seed))
      self.assertAllEqual(y_.shape[:-2], batch_shape)
  
      # Next check that the broadcasting works as expected, and the batch log_prob
      # actually matches the log probs of independent models.
      individual_ssms = []
      for i in range(batch_shape[0]):
          for j in range(batch_shape[1]):
              individual_ssms.append(
                  ConstrainedSeasonalStateSpaceModel(
                      num_timesteps=9,
                      num_seasons=num_seasons,
                      num_steps_per_season=2,
                      drift_scale=drift_scale[i, j, ...],
                      observation_noise_scale=observation_noise_scale[
                          j, ...
                      ],
                      initial_state_prior=mvn_diag.MultivariateNormalDiag(
                          scale_diag=initial_state_prior.scale.diag[
                              j, ...
                          ]
                      ),
                  )
              )
  
      batch_lps_ = self.evaluate(ssm.log_prob(y_)).flatten()
  
      individual_ys = [
          y_[i, j, ...]  # pylint: disable=g-complex-comprehension
          for i in range(batch_shape[0])
          for j in range(batch_shape[1])
      ]
      individual_lps_ = self.evaluate(
          [
              individual_ssm.log_prob(individual_y)
              for (individual_ssm, individual_y) in zip(
                  individual_ssms, individual_ys
              )
          ]
      )
  
      self.assertAllClose(individual_lps_, batch_lps_)
  
imports: import math
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import identity
  from tensorflow_probability.python.bijectors import invert
  from tensorflow_probability.python.bijectors import square
  from tensorflow_probability.python.distributions import categorical
  from tensorflow_probability.python.distributions import inverse_gamma
  from tensorflow_probability.python.distributions import linear_gaussian_ssm
  from tensorflow_probability.python.distributions import lognormal
  from tensorflow_probability.python.distributions import mixture_same_family
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.sts.components.seasonal import ConstrainedSeasonalStateSpaceModel
  from tensorflow_probability.python.sts.components.seasonal import Seasonal
  from tensorflow_probability.python.sts.components.seasonal import SeasonalStateSpaceModel
label_desc: Use image node_id 1 for calling the _ConstrainedSeasonalStateSpaceModelTest obj's underlying member method code with example usage: obj.test_batch_shape() without return types