code: def test_one_step_predictive_with_batch_shape(self):
      num_param_samples = 5
      num_timesteps = 4
      batch_shape = [3, 2]
      observed_time_series = self._build_tensor(
          np.random.randn(*(batch_shape + [num_timesteps]))
      )
      model = self._build_model(
          observed_time_series, prior_batch_shape=batch_shape[1:]
      )
      prior_samples = [
          param.prior.sample(
              num_param_samples, seed=test_util.test_seed()
          )
          for param in model.parameters
      ]
  
      @tf.function(autograph=False, jit_compile=tf.executing_eagerly())
      def _run():
          d = forecast.one_step_predictive(
              model,
              observed_time_series,
              timesteps_are_event_shape=False,
              parameter_samples=prior_samples,
          )
          d_mean = d.mean()
          return d, d_mean, d.log_prob(d_mean)
  
      onestep_dist, onestep_mean, onestep_mean_log_prob = _run()
  
      self.evaluate(tf1.global_variables_initializer())
      self.assertAllEqual(
          onestep_dist.batch_shape_tensor(),
          batch_shape + [num_timesteps],
      )
      self.assertAllEqual(
          tf.shape(onestep_mean), batch_shape + [num_timesteps]
      )
      self.assertAllEqual(
          tf.shape(onestep_mean_log_prob), batch_shape + [num_timesteps]
      )
  
imports: import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.sts import fitting
  from tensorflow_probability.python.sts import forecast
  from tensorflow_probability.python.sts.components import seasonal
  from tensorflow_probability.python.sts.components import sum
  from tensorflow_probability.python.sts.internal import missing_values_util
label_desc: Use image node_id 3 for calling the _ForecastTest obj's underlying member method code with example usage: obj.test_one_step_predictive_with_batch_shape() and returns: d,  d_mean,  d