Update docstring on return type of jvp and vjp (#51035)

Summary:
Updates the docstrings, that `jvp` and `vjp` both return the primal `func_output` first as part of the return tuple,
in line with the docstrings of [hvp](https://github.com/niklasschmitz/pytorch/blob/c620572a3477143df33128318dc0d7d10fab811d/torch/autograd/functional.py#L671) and [vhp](https://github.com/niklasschmitz/pytorch/blob/c620572a3477143df33128318dc0d7d10fab811d/torch/autograd/functional.py#L583).

Pull Request resolved: #51035

Reviewed By: bdhirsh

Differential Revision: D26047693

Pulled By: albanD

fbshipit-source-id: 5f2957a858826b4c1884590b6be7a8bed0791efd

torch/autograd/functional.py

@@ -214,8 +214,11 @@ def vjp(func, inputs, v=None, create_graph=False, strict=False):
             Defaults to ``False``.
 
     Returns:
-        vjp (tuple of Tensors or Tensor): result of the dot product with
-        the same shape as the inputs.
+        output (tuple): tuple with:
+            func_output (tuple of Tensors or Tensor): output of ``func(inputs)``
+
+            vjp (tuple of Tensors or Tensor): result of the dot product with
+            the same shape as the inputs.
 
     Example:
 
@@ -298,8 +301,11 @@ def jvp(func, inputs, v=None, create_graph=False, strict=False):
             Defaults to ``False``.
 
     Returns:
-        jvp (tuple of Tensors or Tensor): result of the dot product with
-        the same shape as the output.
+        output (tuple): tuple with:
+            func_output (tuple of Tensors or Tensor): output of ``func(inputs)``
+
+            jvp (tuple of Tensors or Tensor): result of the dot product with
+            the same shape as the output.
 
     Example: