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

[dynamo] simplify polyfill registration for builtins.all and builtins.any (#133769)

Browse Source 
Pull Request resolved: https://github.com/pytorch/pytorch/pull/133769
Approved by: https://github.com/jansel
This commit is contained in:
Xuehai Pan
2024-08-30 00:47:02 +08:00
committed by PyTorch MergeBot
parent b977abd5de
commit e09324e7da
8 changed files with 51 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
torch/compiler/__init__.py
View File

@ -123,6 +123,7 @@ def allow_in_graph(fn):
def substitute_in_graph(
original_fn: _F,
*,
can_constant_fold_through: bool = False,
skip_signature_check: bool = False,
) -> Callable[[_F], _F]:
"""
@ -142,6 +143,10 @@ def substitute_in_graph(
Args:
original_fn (callable): The original function, usually a C function, to register a polyfill
handler for.
can_constant_fold_through (bool, optional): Whether the polyfill handler can be constant
folded through. That is, if the polyfill handler is a pure function and its arguments
are constant, the result of the polyfill handler can be constant folded during the
compilation. Defaults to ``False``.
skip_signature_check (bool, optional): Whether to skip the signature check between the
original function and the polyfill handler. Defaults to ``False``.
@ -173,6 +178,7 @@ def substitute_in_graph(
return torch._dynamo.substitute_in_graph(
original_fn,
can_constant_fold_through=can_constant_fold_through,
skip_signature_check=skip_signature_check,
)