[dynamo] Remove workaround for functools.wraps in functorch (#142014)

StrongerXi · pytorchmergebot · commit ed77901ec521 · 2024-12-04T19:10:46.000Z
This is no longer needed after #142000. Fixes #123365. Pull Request resolved: #142014 Approved by: https://github.com/zou3519 ghstack dependencies: #142000
diff --git a/torch/_functorch/eager_transforms.py b/torch/_functorch/eager_transforms.py
@@ -599,6 +599,7 @@ def jacrev(
     if not (chunk_size is None or chunk_size > 0):
         raise ValueError("jacrev: `chunk_size` should be greater than 0.")
 
+    @wraps(func)
     def wrapper_fn(*args):
         error_if_complex("jacrev", args, is_input=True)
         vjp_out = _vjp_with_argnums(func, *args, argnums=argnums, has_aux=has_aux)
@@ -761,12 +762,6 @@ def compute_jacobian_preallocate_and_copy():
             return output_input, aux
         return output_input
 
-    # Dynamo does not support HOP composition if their inner function is
-    # annotated with @functools.wraps(...). We circumvent this issue by applying
-    # wraps only if we're not tracing with dynamo.
-    if not torch._dynamo.is_compiling():
-        wrapper_fn = wraps(func)(wrapper_fn)
-
     return wrapper_fn
 
 
@@ -1287,6 +1282,7 @@ def jacfwd(
 
     """
 
+    @wraps(func)
     def wrapper_fn(*args):
         error_if_complex("jacfwd", args, is_input=True)
         primals = args if argnums is None else _slice_argnums(args, argnums)
@@ -1341,12 +1337,6 @@ def push_jvp(basis):
             return tree_unflatten(jac_outs_ins, spec), aux
         return tree_unflatten(jac_outs_ins, spec)
 
-    # Dynamo does not support HOP composition if their inner function is
-    # annotated with @functools.wraps(...). We circumvent this issue by applying
-    # wraps only if we're not tracing with dynamo.
-    if not torch._dynamo.is_compiling():
-        wrapper_fn = wraps(func)(wrapper_fn)
-
     return wrapper_fn
 
 
