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

shape guards (#161178)

Browse Source 
Summary: This PR introduces shape guards to export. Previously only value ranges,  equalities, and specializations would be tracked for symbolic expressions, and we had a forward hook to check them. Instead now we create a function to check shape guards and call it in the exported program.

Test Plan:
updated several tests

Rollback Plan:

Differential Revision: D80713603

Pull Request resolved: https://github.com/pytorch/pytorch/pull/161178
Approved by: https://github.com/tugsbayasgalan
This commit is contained in:
Avik Chaudhuri
2025-09-08 22:44:05 +00:00
committed by PyTorch MergeBot
parent 2c538c9acf
commit 711c8c821e
29 changed files with 617 additions and 123 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
torch/export/__init__.py
View File

@ -540,6 +540,7 @@ def draft_export(
dynamic_shapes: Optional[Union[dict[str, Any], tuple[Any, ...], list[Any]]] = None,
preserve_module_call_signature: tuple[str, ...] = (),
strict: bool = False,
prefer_deferred_runtime_asserts_over_guards: bool = False,
) -> ExportedProgram:
"""
A version of torch.export.export which is designed to consistently produce
@ -555,6 +556,7 @@ def draft_export(
dynamic_shapes=dynamic_shapes,
preserve_module_call_signature=preserve_module_call_signature,
strict=strict,
prefer_deferred_runtime_asserts_over_guards=prefer_deferred_runtime_asserts_over_guards,
)