code: def grad(func: Callable, argnums: Union[int, Tuple[int]] = 0):
      def grad_fn(x_in):
          return ivy.to_native(func(x_in))
  
      def callback_fn(x_in):
          return ivy.to_ivy(
              jax.grad(grad_fn, argnums)(ivy.to_native(x_in))
          )
  
      return callback_fn
  
file_path: ivy/ivy/functional/backends/jax/gradients.py
imports: import jax
  import jax.lax
  from ivy.functional.backends.jax import JaxArray, NativeArray
  from typing import Optional, Callable, Sequence, Union, Tuple
  import ivy
  from ivy.functional.ivy.gradients import _get_required_float_variables, _get_y_and_ret_idxs, _get_native_variables_and_indices, _set_duplicates, _process_func_ret_and_grads
label_desc: Use image node_id 9 for calling a global function with example usage: grad(func, argnums) and returns: callback_fn, ivy, ivy