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



from torch.func import grad
x = torch.randn([])
cos_x = grad(lambda x: torch.sin(x))(x)
assert torch.allclose(cos_x, x.cos())
Second-order gradients
neg_sin_x = grad(grad(lambda x: torch.sin(x)))(x)
assert torch.allclose(neg_sin_x, -x.sin())



When composed with "vmap", "grad" can be used to compute per-
   sample-gradients:



from torch.func import grad, vmap
batch_size, feature_size = 3, 5
def model(weights, feature_vec):
    # Very simple linear model with activation
    assert feature_vec.dim() == 1


