code: def get_state_from_vec(state_vec, state_shape):
      """Inverse of `get_state_vec`."""
      state_sizes = tf.nest.map_structure(ps.reduce_prod, state_shape)
      state_vec_parts = tf.nest.pack_sequence_as(
          state_shape,
          tf.split(state_vec, tf.nest.flatten(state_sizes), axis=-1),
      )
      batch_shape = ps.shape(state_vec)[:-1]
  
      return tf.nest.map_structure(
          lambda sv, s: tf.reshape(
              sv, ps.concat([batch_shape, s], axis=0)
          ),
          state_vec_parts,
          state_shape,
      )
  
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 9 for calling a global function with example usage: get_state_from_vec(state_vec, state_shape) and returns: tf