frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 12:54:11 +08:00
Code Issues Packages Projects Releases Wiki Activity

[dynamo] Better error for invalid @contextlib.contextmanager usage (#156924)

Browse Source 
Fixes #156716

Pull Request resolved: https://github.com/pytorch/pytorch/pull/156924
Approved by: https://github.com/williamwen42
This commit is contained in:
Jason Ansel
2025-06-26 20:18:44 -07:00
committed by PyTorch MergeBot
parent ff8b53c056
commit 60abb0d327
2 changed files with 22 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
torch/_dynamo/graph_break_registry.json
View File

@ -2177,5 +2177,16 @@
"INFO"
]
}
],
"GB0221": [
{
"Gb_type": "non-generator contextlib.contextmanager",
"Context": "str(self.vt.get_code())",
"Explanation": "Cannot compile function decorated with `@contextlib.contextmanager` that is not a generator, i.e. does not use `yield`",
"Hints": [
"Use `yield` in the function body instead of `return`.",
"Remove the `@contextlib.contextmanager` decorator."
]
}
]
}

12
torch/_dynamo/variables/functions.py
View File

@ -942,7 +942,17 @@ class LocalGeneratorFunctionVariable(BaseUserFunctionVariable):
args: "list[VariableTracker]",
kwargs: "dict[str, VariableTracker]",
) -> "VariableTracker":
assert is_generator(self.vt.get_code())
if not is_generator(self.vt.get_code()):
unimplemented_v2(
gb_type="non-generator contextlib.contextmanager",
context=str(self.vt.get_code()),
explanation="Cannot compile function decorated with `@contextlib.contextmanager` that is not a generator"
", i.e. does not use `yield`",
hints=[
"Use `yield` in the function body instead of `return`.",
"Remove the `@contextlib.contextmanager` decorator.",
],
)
inline_tracer = self._build_inline_tracer(tx, args, kwargs)
code = self.vt.get_code()