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

Enhance torch.vmap support from inside torch.compile (#116050)

Browse Source 
This work rewrites vmap support in torch.compile by inlining most of
the frames into the existing FX graph. It also unlocks to PyTorch to
support features that were previously missing, such as keyword args.

Fixes: https://github.com/pytorch/pytorch/issues/114306

Pull Request resolved: https://github.com/pytorch/pytorch/pull/116050
Approved by: https://github.com/zou3519
This commit is contained in:
Guilherme Leobas
2024-01-19 12:36:51 -03:00
committed by PyTorch MergeBot
parent b2a3d6ba0d
commit 80cf0ce153
18 changed files with 642 additions and 412 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
.flake8
View File

@ -34,6 +34,8 @@ per-file-ignores =
torch/utils/cpp_extension.py: B950
torchgen/api/types/__init__.py: F401,F403
torchgen/executorch/api/types/__init__.py: F401,F403
test/dynamo/test_higher_order_ops.py: B950
torch/testing/_internal/dynamo_test_failures.py: B950
optional-ascii-coding = True
exclude =
./.git,