code: def grad(func, argnum=None):
      
      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
docs: 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)