code: def grad(func, argnum=None):
      """Return function that computes gradient of arguments.
  
      Parameters
      ----------
      func: a python function
          The forward (loss) function.
      argnum: an int or a list of int
          The index of argument to calculate gradient for.
  
      Returns
      -------
      grad_func: a python function
          A function that would compute the gradient of arguments.
  
      Examples
      --------
      >>> # autograd supports dynamic graph which is changed
      >>> # every instance
      >>> def func(x):
      >>>     r = random.randint(0, 1)
      >>>     if r % 2:
      >>>         return x**2
      >>>     else:
      >>>         return x/3
      >>> # use `grad(func)` to get the gradient function
      >>> for x in range(10):
      >>>     grad_func = grad(func)
      >>>     inputs = nd.array([[1, 2, 3], [4, 5, 6]])
      >>>     grad_vals = grad_func(inputs)
      """
      grad_with_loss_func = grad_and_loss(func, argnum)
  
      @functools.wraps(grad_with_loss_func)
      def wrapped(*args):
          return grad_with_loss_func(*args)[0]
  
      return wrapped
  
imports: import functools
  import mxnet.ndarray
  from mxnet.ndarray import zeros_like
  from mxnet.autograd import *
  from mxnet.test_utils import *
  from common import xfail_when_nonstandard_decimal_separator
  from mxnet.test_utils import environment
  import pytest
label_desc: Use image node_id 2 for calling a global function with example usage: grad(func, argnum) and returns: wrapped, grad_with_loss_func