the function is a tuple of the gradient tensor(s) with respect
      to the first argument and the primal computation. If specified
      "has_aux" equals "True", tuple of gradients and tuple of the
      forward computation with output auxiliary objects is returned.
      If "argnums" is a tuple of integers, a tuple of a tuple of the
      output gradients with respect to each "argnums" value and the
      forward computation is returned.
Return type:
      Callable
See "grad()" for examples