OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-07 12:21:27 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity

Unconditionally enable python dispatcher in AOTAutograd (#88365)

Browse Source 
Signed-off-by: Edward Z. Yang <ezyang@fb.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/88365
Approved by: https://github.com/Chillee
This commit is contained in:
Edward Z. Yang 2022-11-02 18:55:33 -07:00 committed by PyTorch MergeBot
parent a689502275
commit 97d3b200ca
2 changed files with 17 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
functorch/_src/aot_autograd.py
View File

@ -391,6 +391,7 @@ def aot_dispatch_autograd(flat_fn, flat_args: List[Tensor], aot_config: AOTConfi
disable_amp = torch._C._is_any_autocast_enabled() disable_amp = torch._C._is_any_autocast_enabled()
if config.use_functionalize: if config.use_functionalize:
with enable_python_dispatcher():
# Trace once without decompositions, into a graph of ATen ops. # Trace once without decompositions, into a graph of ATen ops.
# NB: tracing_mode is real, as it's assumed the calling context setup # NB: tracing_mode is real, as it's assumed the calling context setup
# fake tensor mode / symbolic shapes if that is needed # fake tensor mode / symbolic shapes if that is needed

2
torch/_ops.py
View File

@ -103,7 +103,7 @@ def resolve_key(op: PyOperatorABC, k: DispatchKey): # type: ignore[valid-type]
# The dispatch key itself will implicitly route to backend fallback. # The dispatch key itself will implicitly route to backend fallback.
# This is probably not great for the pure Python implementation. # This is probably not great for the pure Python implementation.
return k return k
raise RuntimeError("could not find kernel") raise NotImplementedError(f"could not find kernel for {op} at dispatch key {k}")
pyop_namespace = {} pyop_namespace = {}