code: def _flatten_and_concat(x, batch_shape, dtype):
      """Flattens and concatenates a structured `x`."""
      # For convenience.
      if x is None:
          return x
  
      def _reshape_part(part):
          part = tf.cast(part, dtype)
          new_shape = ps.concat(
              [batch_shape, [-1]],
              axis=-1,
          )
          return tf.reshape(part, ps.cast(new_shape, tf.int32))
  
      x = tf.nest.map_structure(_reshape_part, x)
      return tf.concat(tf.nest.flatten(x), axis=-1)
  
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 1 for calling a global function with example usage: _flatten_and_concat(x, batch_shape, dtype) and returns: tf, x, tf