code: def batch_concat_selection(
      observation_dict: Dict[str, types.NestedTensor],
      concat_keys: Optional[Iterable[str]] = None,
      output_dtype=tf.float32,
  ) -> tf.Tensor:
      """Concatenate a dict of observations into 2-D tensors."""
      concat_keys = concat_keys or sorted(observation_dict.keys())
      to_concat = []
      for obs in concat_keys:
          if obs not in observation_dict:
              raise KeyError(
                  "Missing observation. Requested: {} (available: {})".format(
                      obs, list(observation_dict.keys())
                  )
              )
          to_concat.append(tf.cast(observation_dict[obs], output_dtype))
  
      return tf2_utils.batch_concat(to_concat)
  
file_path: acme/acme/agents/tf/svg0_prior/utils.py
imports: import collections
  from typing import Tuple, Optional, Dict, Iterable
  from acme import types
  from acme.tf import utils
  import sonnet
  import tensorflow
  import tree
label_desc: Use image node_id 3 for calling a global function with example usage: batch_concat_selection(observation_dict, concat_keys, output_dtype) and returns: tf2_utils