code: def test_multiple_chains_example(self):
      batch_shape = [2, 3]
      num_timesteps = 5
      num_results = 6
      num_chains = 4
  
      # Use an observation mask to additionally test that masks are
      # threaded through the HMC (and VI) APIs.
      observed_time_series_ = np.random.randn(
          *(batch_shape + [num_timesteps])
      )
      observed_time_series = missing_values_util.MaskedTimeSeries(
          self._build_tensor(observed_time_series_),
          is_missing=self._build_tensor(
              [False, True, False, False, True], dtype=np.bool_
          ),
      )
  
      model = self._build_model(observed_time_series)
      samples, kernel_results = fitting.fit_with_hmc(
          model,
          observed_time_series,
          num_results=num_results,
          chain_batch_shape=num_chains,
          num_warmup_steps=2,
          num_variational_steps=2,
      )
  
      self.evaluate(tf1.global_variables_initializer())
      samples_, kernel_results_ = self.evaluate(
          (samples, kernel_results)
      )
  
      acceptance_rate = np.mean(
          kernel_results_.inner_results.inner_results.is_accepted,
          axis=0,
      )
  
      # Combining the samples from multiple chains into a single dimension allows
      # us to easily pass sampled parameters to downstream forecasting methods.
      combined_samples_ = [
          np.reshape(param_draws, [-1] + list(param_draws.shape[2:]))
          for param_draws in samples_
      ]
  
      self.assertAllEqual(
          acceptance_rate.shape, [num_chains] + batch_shape
      )
      for parameter, samples_ in zip(
          model.parameters, combined_samples_
      ):
          self.assertAllEqual(
              samples_.shape,
              [num_results * num_chains]
              + self._batch_shape_as_list(parameter.prior)
              + self._event_shape_as_list(parameter.prior),
          )
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.sts import fitting
  from tensorflow_probability.python.sts.components import local_linear_trend
  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
  from tensorflow_probability.python.vi import csiszar_divergence
  from tensorflow_probability.python.vi import optimization
label_desc: Use image node_id 3 for calling the _HMCTests obj's underlying member method code with example usage: obj.test_multiple_chains_example() without return types