"(output, aux)" tuple where the first element is the output of
        the function to be differentiated and the second element is
        other auxiliary objects that will not be differentiated.
        Default: False.
Returns:
      Returns a "(output, jvp_out)" tuple containing the output of
      "func" evaluated at "primals" and the Jacobian-vector product.
      If "has_aux is True", then instead returns a "(output, jvp_out,
      aux)" tuple.
Note:
 You may see this API error out with "forward-mode AD not
 implemented for operator X". If so, please file a bug report and
 we will prioritize it.

jvp is useful when you wish to compute gradients of a function R^1
   -> R^N



from torch.func import jvp
x = torch.randn([])
f = lambda x: x * torch.tensor([1., 2., 3])
value, grad = jvp(f, (x,), (torch.tensor(1.),))
assert torch.allclose(value, f(x))
assert torch.allclose(grad, torch.tensor([1., 2, 3]))


