code: def get_ode_fn_vec(ode_fn, state_shape):
      """Returns a wrapper around the user-specified `ode_fn` argument.
  
      The second argument and output of `ode_fn(time, state)` are N-D `Tensor`s.
      This function returns a wrapper `ode_fn_vec(time, state_vec)` whose
      second argument and output are 1-D `Tensor`s corresponding to reshaped
      versions of `state` and `ode_fn(time, state)`.
  
      Args:
        ode_fn: User-specified `ode_fn` passed to `solve`.
        state_shape: The shape of the second argument and output of `ode_fn`.
  
      Returns:
        The wrapper described above.
      """
  
      def ode_fn_vec(time, state_vec):
          return get_state_vec(
              ode_fn(time, get_state_from_vec(state_vec, state_shape))
          )
  
      return ode_fn_vec
  
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 10 for calling a global function with example usage: get_ode_fn_vec(ode_fn, state_shape) and returns: ode_fn_vec, get_state_vec