Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 3 commits into from
Closed

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

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
f064b0d
Update
StrongerXi Dec 3, 2024
572aa21
Update
StrongerXi Dec 4, 2024
abf8368
Update on "[dynamo] Remove workaround for `functools.wraps` in functo…
StrongerXi Dec 5, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 2 additions & 12 deletions torch/_functorch/eager_transforms.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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


Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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


Expand Down
Loading