Skip to content

Disbale torchdynamo on AOT Autograd generated graphs #662

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

Merged
merged 1 commit into from
Apr 6, 2022
Merged

Disbale torchdynamo on AOT Autograd generated graphs #662

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
8 changes: 8 additions & 0 deletions functorch/_src/aot_autograd.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,12 @@
from .named_members_polyfill import _named_parameters, _named_buffers
from typing import Callable, List, Dict, Any, Tuple, Optional

try:
from torchdynamo import disable as disable_torchdynamo
except ImportError:
def disable_torchdynamo(x):
return x

pytree._register_pytree_node(
immutable_collections.immutable_list,
lambda x: (list(x), None),
Expand Down Expand Up @@ -129,6 +135,7 @@ def create_aot_autograd_function(

class CompiledFunction(torch.autograd.Function):
@staticmethod
@disable_torchdynamo
def forward(ctx, *flat_tensor_args):
nonlocal compiled_fw, compiled_bw, num_outs
if compiled_fw is None:
Expand Down Expand Up @@ -163,6 +170,7 @@ def forward(ctx, *flat_tensor_args):
return tuple(fw_outs[0:num_outs])

@staticmethod
@disable_torchdynamo
def backward(ctx, *flat_args):
contiguous_args = [t.contiguous() for t in flat_args]
# contiguous_args = [t for t in flat_args]
Expand Down