code: def testBatchShape(self):
      # Check that the model builds with batches of parameters.
      order = 3
      batch_shape = [4, 2]
  
      # No `_build_placeholder`, because coefficients must have static shape.
      coefficients = np.random.randn(*(batch_shape + [order])).astype(
          self.dtype
      )
      order = max(
          order, order + 1
      )  # shape of initial_state_prior, scale_diag
  
      level_scale = self._build_placeholder(
          np.exp(np.random.randn(*batch_shape))
      )
  
      ssm = AutoregressiveMovingAverageStateSpaceModel(
          num_timesteps=10,
          ar_coefficients=coefficients,
          ma_coefficients=coefficients,
          level_scale=level_scale,
          initial_state_prior=mvn_diag.MultivariateNormalDiag(
              scale_diag=self._build_placeholder(np.ones([order]))
          ),
      )
      if self.use_static_shape:
          self.assertAllEqual(
              tensorshape_util.as_list(ssm.batch_shape), batch_shape
          )
      else:
          self.assertAllEqual(ssm.batch_shape_tensor(), batch_shape)
  
      y = ssm.sample(seed=test_util.test_seed(sampler_type="stateless"))
      if self.use_static_shape:
          self.assertAllEqual(
              tensorshape_util.as_list(y.shape)[:-2], batch_shape
          )
      else:
          self.assertAllEqual(tf.shape(y)[:-2], batch_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 tensorshape_util
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.sts.components.autoregressive import AutoregressiveStateSpaceModel
  from tensorflow_probability.python.sts.components.autoregressive_moving_average import AutoregressiveMovingAverageStateSpaceModel
label_desc: Use image node_id 3 for calling the _AutoregressiveMovingAverageStateSpaceModelTest obj's underlying member method code with example usage: obj.testBatchShape() without return types