code: def _make_vector_sample_transformations(
      model, flatten_sample_transformations
  ):
      """Makes `model`'s sample transformations compatible with vector events."""
      sample_transformations = collections.OrderedDict()
  
      def flattened_transform(x, dtype, transform):
          res = transform(_split_and_reshape_event(x, model))
          if not flatten_sample_transformations:
              return res
          batch_shape = ps.shape(x)[:-1]
          return _flatten_and_concat(res, batch_shape, dtype)
  
      # We yank this out to avoid capturing the loop variable.
      def make_flattened_transform(transform_name, transform):
          if flatten_sample_transformations:
              dtype = _get_unique_dtype(
                  transform.dtype,
                  lambda d: "Sample transformation '{}' must have only one Tensor dtype, "  # pylint: disable=g-long-lambda
                  "saw: {}".format(transform_name, d),
              )
              transform = transform._replace(
                  ground_truth_mean=_flatten_and_concat(
                      transform.ground_truth_mean, (), dtype
                  ),
                  ground_truth_standard_deviation=_flatten_and_concat(
                      transform.ground_truth_standard_deviation,
                      (),
                      dtype,
                  ),
                  ground_truth_mean_standard_error=_flatten_and_concat(
                      transform.ground_truth_mean_standard_error,
                      (),
                      dtype,
                  ),
                  ground_truth_standard_deviation_standard_error=_flatten_and_concat(
                      transform.ground_truth_standard_deviation_standard_error,
                      (),
                      dtype,
                  ),
              )
          else:
              dtype = None
          return transform._replace(
              fn=functools.partial(
                  flattened_transform, dtype=dtype, transform=transform
              )
          )
  
      for key, transform in model.sample_transformations.items():
          sample_transformations[key] = make_flattened_transform(
              key, transform
          )
      return sample_transformations
  
file_path: probability/spinoffs/inference_gym/inference_gym/targets/vector_model.py
imports: import collections
  import functools
  import tensorflow.compat.v2
  import tensorflow_probability
  from tensorflow_probability.python.internal import prefer_static
  from inference_gym.targets import model
  from tensorflow.python.util import nest
label_desc: Use image node_id 6 for calling a global function with example usage: _make_vector_sample_transformations(model, flatten_sample_transformations) and returns: sample_transformations, _flatten_and_concat, transform, res