code: def get_jacobian_fn_mat(jacobian_fn, ode_fn_vec, state_shape, dtype):
      """Returns a wrapper around the user-specified `jacobian_fn` argument.
  
      `jacobian_fn` is an optional argument that can either be a constant `Tensor`
      or a function of the form `jacobian_fn(time, state)`. This function returns a
      wrapper `jacobian_fn_mat(time, state_vec)` whose second argument and output
      are 1 and 2-D `Tensor`s, respectively, corresponding reshaped versions of
      `state` and `jacobian_fn(time, state)`.
  
      Args:
        jacobian_fn: User-specified `jacobian_fn` passed to `solve`.
        ode_fn_vec: Result of `get_ode_fn_vec`.
        state_shape: The shape of the second argument and output of `ode_fn`.
        dtype: If `jacobian_fn` is constant, what dtype to convert it to.
  
      Returns:
        The wrapper described above.
      """
      if jacobian_fn is None:
          return _AutomaticJacobian(ode_fn_vec)
  
      if not callable(jacobian_fn):
          jacobian_fn = tf.nest.map_structure(
              lambda x: tf.convert_to_tensor(x, dtype=dtype),
              jacobian_fn,
          )
          constant_jacobian_mat = _flatten_nested_jacobian(
              jacobian_fn, state_shape
          )
  
          def constant_jacobian_fn_mat(*_):
              return constant_jacobian_mat
  
          return constant_jacobian_fn_mat
  
      def jacobian_fn_mat(time, state_vec):
          return _flatten_nested_jacobian(
              jacobian_fn(
                  time, get_state_from_vec(state_vec, state_shape)
              ),
              state_shape,
          )
  
      return jacobian_fn_mat
  
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 7 for calling a global function with example usage: get_jacobian_fn_mat(jacobian_fn, ode_fn_vec, state_shape, dtype) and returns: jacobian_fn_mat, _AutomaticJacobian, constant_jacobian_fn_mat, _flatten_nested_jacobian, constant_jacobian_mat