code: def pad_batch_dimension_for_multiple_chains(
      observed_time_series, model, chain_batch_shape
  ):
      """ "Expand the observed time series with extra batch dimension(s)."""
      # Running with multiple chains introduces an extra batch dimension. In
      # general we also need to pad the observed time series with a matching batch
      # dimension.
      #
      # For example, suppose our model has batch shape [3, 4] and
      # the observed time series has shape `concat([[5], [3, 4], [100])`,
      # corresponding to `sample_shape`, `batch_shape`, and `num_timesteps`
      # respectively. The model will produce distributions with batch shape
      # `concat([chain_batch_shape, [3, 4]])`, so we pad `observed_time_series` to
      # have matching shape `[5, 1, 3, 4, 100]`, where the added `1` dimension
      # between the sample and batch shapes will broadcast to `chain_batch_shape`.
  
      [  # Extract mask and guarantee `event_ndims=2`.
          observed_time_series,
          is_missing,
      ] = canonicalize_observed_time_series_with_mask(
          observed_time_series
      )
  
      event_ndims = (
          2  # event_shape = [num_timesteps, observation_size=1]
      )
  
      model_batch_ndims = (
          tensorshape_util.rank(model.batch_shape)
          if tensorshape_util.rank(model.batch_shape) is not None
          else tf.shape(model.batch_shape_tensor())[0]
      )
  
      # Compute ndims from chain_batch_shape.
      chain_batch_shape = tf.convert_to_tensor(
          value=chain_batch_shape,
          name="chain_batch_shape",
          dtype=tf.int32,
      )
      if not tensorshape_util.is_fully_defined(chain_batch_shape.shape):
          raise ValueError(
              "Batch shape must have static rank. (given: {})".format(
                  chain_batch_shape
              )
          )
      if tensorshape_util.rank(chain_batch_shape.shape) == 0:
          # expand int `k` to `[k]`.
          chain_batch_shape = chain_batch_shape[tf.newaxis]
      chain_batch_ndims = tf.compat.dimension_value(
          chain_batch_shape.shape[0]
      )
  
      def do_padding(observed_time_series_tensor):
          current_sample_shape = ps.shape(observed_time_series_tensor)[
              : -(model_batch_ndims + event_ndims)
          ]
          current_batch_and_event_shape = ps.shape(
              observed_time_series_tensor
          )[-(model_batch_ndims + event_ndims) :]
          return tf.reshape(
              tensor=observed_time_series_tensor,
              shape=ps.concat(
                  [
                      current_sample_shape,
                      ps.ones([chain_batch_ndims], dtype=tf.int32),
                      current_batch_and_event_shape,
                  ],
                  axis=0,
              ),
          )
  
      # Padding is only needed if the observed time series has sample shape.
      observed_time_series = ps.cond(
          ps.rank(observed_time_series)
          > model_batch_ndims + event_ndims,
          lambda: do_padding(observed_time_series),
          lambda: observed_time_series,
      )
      if is_missing is not None:
          is_missing = ps.cond(
              ps.rank(is_missing) > model_batch_ndims + event_ndims,
              lambda: do_padding(is_missing),
              lambda: is_missing,
          )
      return missing_values_util.MaskedTimeSeries(
          observed_time_series, is_missing=is_missing
      )
  
imports: import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import categorical
  from tensorflow_probability.python.distributions import mixture_same_family
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import mvn_linear_operator
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.sts.internal import missing_values_util
label_desc: Use image node_id 3 for calling a global function with example usage: pad_batch_dimension_for_multiple_chains(observed_time_series, model, chain_batch_shape) and returns: missing_values_util, tf