[functorch] introduce an experimental map() op.#88767
[functorch] introduce an experimental map() op.#88767zhxchen17 wants to merge 1 commit intopytorch:masterfrom
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/88767
Note: Links to docs will display an error until the docs builds have been completed. ❗ 1 Active SEVsThere are 1 currently active SEVs. If your PR is affected, please view them below: ✅ No FailuresAs of commit 4017a0b: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
This pull request was exported from Phabricator. Differential Revision: D41165796 |
|
What exactly is being looped over here? Are you permitting lists of varying sizes in the IR format? (If so, that makes me very negative on this feature. cond is OK because it doesn't introduce new data types to the FX IR. This would.) |
I think the semantics is very similar to how jax.lax.map works, basically we loop over the second argument:
No we don't permit that, varying size list is not useful to us and we don't want to add it. |
bfdb62e to
1fc1088
Compare
|
This pull request was exported from Phabricator. Differential Revision: D41165796 |
1 similar comment
|
This pull request was exported from Phabricator. Differential Revision: D41165796 |
1fc1088 to
ec4bfb2
Compare
ec4bfb2 to
331dbe5
Compare
|
This pull request was exported from Phabricator. Differential Revision: D41165796 |
|
OK, makes more sense. |
|
deferring mostly to @voznesenskym and @zou3519 for review on this. Holler if you want me to look at something specific. |
zou3519
left a comment
There was a problem hiding this comment.
I have some questions, but this seems reasonable to me, it's in line with what we did with cond
|
@zhxchen17 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
1fc1088 to
e4c39a0
Compare
|
addressed all comments. |
b81d463 to
631790b
Compare
Summary: I saw the following issue only on Windows build in PR #88767: ``` RuntimeError: AttributeError: 'SymNode' object has no attribute 'torch::impl::PythonSymNodeImpl::ge' ``` The reason it's only on Windows is because we get the attributes of SymNode in C++ with `__FUNCTION__` macro, which is not in C++ standard, therefore has platform specific behavior. In this case, MSVC will include a function's namespace and class name, which is not intended here. Instead we should use `__func__`. see: https://en.cppreference.com/w/c/language/function_definition godbolt example to show the difference: https://godbolt.org/z/PGfvecxPx Test Plan: CI Reviewers: Subscribers: Tasks: Tags: ghstack-source-id: 57cbfab Pull Request resolved: #89264
Summary: I saw the following issue only on Windows build in PR #88767: ``` RuntimeError: AttributeError: 'SymNode' object has no attribute 'torch::impl::PythonSymNodeImpl::ge' ``` The reason it's only on Windows is because we get the attributes of SymNode in C++ with `__FUNCTION__` macro, which is not in C++ standard, therefore has platform specific behavior. In this case, MSVC will include a function's namespace and class name, which is not intended here. Instead we should use `__func__`. see: https://en.cppreference.com/w/c/language/function_definition godbolt example to show the difference: https://godbolt.org/z/PGfvecxPx Test Plan: CI Reviewers: Subscribers: Tasks: Tags: [ghstack-poisoned]
Summary: I saw the following issue only on Windows build in PR #88767: ``` RuntimeError: AttributeError: 'SymNode' object has no attribute 'torch::impl::PythonSymNodeImpl::ge' ``` It's only on Windows because we get the attributes of SymNode in C++ with `__FUNCTION__` macro, which is not in C++ standard, therefore has platform specific behavior. In this case, MSVC will include a function's namespace and class name, which is not intended here. Instead we should use `__func__`. see: https://en.cppreference.com/w/cpp/language/function#Function_definition godbolt example to show the difference: https://godbolt.org/z/PGfvecxPx Test Plan: CI Reviewers: Subscribers: Tasks: Tags: Pull Request resolved: #89264 Approved by: https://github.com/ezyang
Summary:
We want to introduce an experimental control flow op: map() to export some models as FX graphs correctly.
Some calrification on basic requirements we have in mind:
1. This op can nest cond() and other control flow primitives internally.
2. We don't necessarily need loop carried dependencies for the models we've seen.
3. This map() op can handle dynamically shaped tensor as input and return dynamically shaped output based on input shapes.
4. We should be able to pass through additional arguments to the loop body as extra arguments.
In this diff we introduce a new control flow op `map()` which has the following semantics:
```
def map(f: Callable, xs: Tensor, *args):
# one possible implementation:
return torch.stack([f(x, *args) for x in xs])
```
Test Plan:
pytest functorch/test_control_flow.py
CI
Reviewers:
Subscribers:
Tasks:
Tags:
631790b to
4017a0b
Compare
|
@zhxchen17 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
…nctorch.experimental.control_flow. Summary: Similar to #88767 we want to reduce the chance that users accidentally import private functions from `functorch.experimental.cond` as if they were public interfaces. We also move `cond()` under `control_flow.py` to stay consistent with `map()` op. Test Plan: CI Reviewers: Subscribers: Tasks: Tags: [ghstack-poisoned]
…nctorch.experimental.control_flow. Summary: Similar to #88767 we want to reduce the chance that users accidentally import private functions from `functorch.experimental.cond` as if they were public interfaces. We also move `cond()` under `control_flow.py` to stay consistent with `map()` op. Test Plan: CI Reviewers: Subscribers: Tasks: Tags: ghstack-source-id: bcce75d Pull Request resolved: #89819
…nctorch.experimental.control_flow. (#89819) Summary: Similar to #88767 we want to reduce the chance that users accidentally import private functions from `functorch.experimental.cond` as if they were public interfaces. We also move `cond()` under `control_flow.py` to stay consistent with `map()` op. Test Plan: CI Reviewers: Subscribers: Tasks: Tags: Pull Request resolved: #89819 Approved by: https://github.com/zou3519
Summary: I saw the following issue only on Windows build in PR pytorch#88767: ``` RuntimeError: AttributeError: 'SymNode' object has no attribute 'torch::impl::PythonSymNodeImpl::ge' ``` It's only on Windows because we get the attributes of SymNode in C++ with `__FUNCTION__` macro, which is not in C++ standard, therefore has platform specific behavior. In this case, MSVC will include a function's namespace and class name, which is not intended here. Instead we should use `__func__`. see: https://en.cppreference.com/w/cpp/language/function#Function_definition godbolt example to show the difference: https://godbolt.org/z/PGfvecxPx Test Plan: CI Reviewers: Subscribers: Tasks: Tags: Pull Request resolved: pytorch#89264 Approved by: https://github.com/ezyang
Summary:
We want to introduce an experimental control flow op: map() to export some models as FX graphs correctly.
Some calrification on basic requirements we have in mind:
1. This op can nest cond() and other control flow primitives internally.
2. We don't necessarily need loop carried dependencies for the models we've seen.
3. This map() op can handle dynamically shaped tensor as input and return dynamically shaped output based on input shapes.
4. We should be able to pass through additional arguments to the loop body as extra arguments.
In this diff we introduce a new control flow op `map()` which has the following semantics:
```
def map(f: Callable, xs: Tensor, *args):
# one possible implementation:
return torch.stack([f(x, *args) for x in xs])
```
Test Plan:
pytest functorch/test_control_flow.py
CI
Differential Revision: D41165796
Pull Request resolved: pytorch#88767
Approved by: https://github.com/zou3519
…nctorch.experimental.control_flow. (pytorch#89819) Summary: Similar to pytorch#88767 we want to reduce the chance that users accidentally import private functions from `functorch.experimental.cond` as if they were public interfaces. We also move `cond()` under `control_flow.py` to stay consistent with `map()` op. Test Plan: CI Reviewers: Subscribers: Tasks: Tags: Pull Request resolved: pytorch#89819 Approved by: https://github.com/zou3519
Summary:
We want to introduce an experimental control flow op: map() to export some models as FX graphs correctly.
Some calrification on basic requirements we have in mind:
In this diff we introduce a new control flow op
map()which has the following semantics:Test Plan:
pytest functorch/test_control_flow.py
CI
Differential Revision: D41165796