code: def get_state_vec(state):
      """Converts a possibly nested state into a vector."""
      return tf.concat(
          tf.nest.flatten(
              tf.nest.map_structure(
                  lambda s: tf.reshape(s, [-1]), state
              )
          ),
          axis=-1,
      )
  
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.math import gradient
  from tensorflow.python.util import nest
label_desc: Use image node_id 8 for calling a global function with example usage: get_state_vec(state) and returns: tf