Enable `in_dims` for vmap frontend api by zou3519 · Pull Request #40717 · pytorch/pytorch

zou3519 · 2020-06-29T19:24:11Z

Stack from ghstack:

Enable in_dims for vmap frontend api #40717 Enable in_dims for vmap frontend api

in_dims specifies which dimension of the input tensors should be
vmapped over. One can also specify None as an in_dim for a particular
input to indicate that we do not map over said input.

We implement in_dims by creating a BatchedTensor with BatchDim equal
to said in_dim. Most of this PR is error checking. in_dims must
satisfy the following:

in_dim can be either an int or a Tuple[Optional[int]]. If it is an
int, we use it to mean the in_dim for every input.
If in_dims is not-None at some index idx, then the input at index
idx MUST be a tensor (vmap can only map over tensors).

jax supports something more generalized: their in_dims can match the
structure of the inputs to the function (i.e., it is a nested python
data structure matching the data structure of inputs specifying where
in inputs the Tensors to be mapped are and what their map dims should
be). We don't have the infrastruture yet so we only support int or a
flat tuple for in_dims.

Test Plan:

pytest test/test_vmap.py -v

Differential Revision: D22397914

`in_dims` specifies which dimension of the input tensors should be vmapped over. One can also specify `None` as an `in_dim` for a particular input to indicate that we do not map over said input. We implement `in_dims` by creating a BatchedTensor with BatchDim equal to said `in_dim`. Most of this PR is error checking. `in_dims` must satisfy the following: - `in_dim` can be either an int or a Tuple[Optional[int]]. If it is an int, we use it to mean the `in_dim` for every input. - If `in_dims` is not-None at some index `idx`, then the input at index `idx` MUST be a tensor (vmap can only map over tensors). jax supports something more generalized: their `in_dims` can match the structure of the `inputs` to the function (i.e., it is a nested python data structure matching the data structure of `inputs` specifying where in `inputs` the Tensors to be mapped are and what their map dims should be). We don't have the infrastruture yet so we only support `int` or a flat tuple for `in_dims`. Test Plan: - `pytest test/test_vmap.py -v` [ghstack-poisoned]

`in_dims` specifies which dimension of the input tensors should be vmapped over. One can also specify `None` as an `in_dim` for a particular input to indicate that we do not map over said input. We implement `in_dims` by creating a BatchedTensor with BatchDim equal to said `in_dim`. Most of this PR is error checking. `in_dims` must satisfy the following: - `in_dim` can be either an int or a Tuple[Optional[int]]. If it is an int, we use it to mean the `in_dim` for every input. - If `in_dims` is not-None at some index `idx`, then the input at index `idx` MUST be a tensor (vmap can only map over tensors). jax supports something more generalized: their `in_dims` can match the structure of the `inputs` to the function (i.e., it is a nested python data structure matching the data structure of `inputs` specifying where in `inputs` the Tensors to be mapped are and what their map dims should be). We don't have the infrastruture yet so we only support `int` or a flat tuple for `in_dims`. Test Plan: - `pytest test/test_vmap.py -v` ghstack-source-id: f77f2dc Pull Request resolved: #40717

dr-ci · 2020-06-29T19:44:38Z

💊 CI failures summary and remediations

As of commit 7259721 (more details on the Dr. CI page):

💚 💚 Looks good so far! There are no failures yet. 💚 💚

This comment was automatically generated by Dr. CI (expand for details).

Follow this link to opt-out of these comments for your Pull Requests.

Please report bugs/suggestions on the GitHub issue tracker or post in the (internal) Dr. CI Users group.

See how this bot performed.

This comment has been revised 7 times.

ezyang · 2020-07-06T14:31:25Z

+            fn=fn_name, in_dims=in_dims, num_inputs=len(args)))
+
+    if len(args) == 0:
+        raise ValueError(NO_INPUTS.format(fn=fn_name))


Is there any reason to block the zero length args case, besides "then vmap doesn't do anything"? I'm thinking of how people have found it useful to do zero-size batches; it may be harmless to have zero length args (unless it is not?)

The only two reasons I have are: (1) "then vmap doesn't do anything" and (2) "jax doesn't allow it". I agree that it seems harmless to have zero-length arguments.

It's not too hard to modify this to work so I'll add this as a follow-up for later (and think more about if it is actually harmless or not)

ezyang · 2020-07-06T14:32:44Z

+EXPECTED_IN_DIMS_TO_BE_INT_OR_TUPLE = (
+    'vmap({fn}, in_dims={in_dims}, ...): expected `in_dims` to be int or tuple, '
+    'got: {actual_type}.'
+)


No action needed on this comment: I personally prefer having messages inline at their use sites, if they're only used once. Makes it easier to see what the error message is and ensure that the format string is up to date :) (also, you can't use f-strings in this style!)

Did we drop support for Python < 3.6? (I know we dropped support for Python 2, but I didn't realize the < 3.6 part)

I agree with your comment, reading 66 lines of error messages at the top of the file and away from the callsites makes me sad. Will fix in a follow-up.

ezyang · 2020-07-06T14:34:48Z


+# Check compatibility of `in_dims` and `args`. More specifically, checks the following:
+# Wherever an in_dim is not None, then the corresponding index in args must be
+# a Tensor. Furthermore, tensor must have the `in_dim` (0 <= in_dim < tensor.dim())


Type signature on this function (and the others) would be very helpful!

RIP Python 2!!!

I added type hints to all of the functions :D. We'll have to relax some of these in the future when we support accepting arbitrary nested python data structures, but these do make the code easier to read now.

ezyang · 2020-07-06T14:39:15Z

+# Wherever an in_dim is not None, then the corresponding index in args must be
+# a Tensor. Furthermore, tensor must have the `in_dim` (0 <= in_dim < tensor.dim())
+def _check_args_can_be_mapped_with_in_dims(in_dims_as_tuple, args, fn_name, in_dims):
+    for idx, (in_dim, arg) in enumerate(zip(in_dims_as_tuple, args)):


If you extend this to work on arbitrary Python collections as opposed to just tuples, zipping here isn't going to work anymore, right? Would we expect in-dims to also have the same "shape" as args, in this case?

That's correct. Extending this to work on arbitrary Python collections would make it so that we need new error validation code here. Furthermore, we'd expect in_dims to have the same "shape" as args.

`in_dims` specifies which dimension of the input tensors should be vmapped over. One can also specify `None` as an `in_dim` for a particular input to indicate that we do not map over said input. We implement `in_dims` by creating a BatchedTensor with BatchDim equal to said `in_dim`. Most of this PR is error checking. `in_dims` must satisfy the following: - `in_dim` can be either an int or a Tuple[Optional[int]]. If it is an int, we use it to mean the `in_dim` for every input. - If `in_dims` is not-None at some index `idx`, then the input at index `idx` MUST be a tensor (vmap can only map over tensors). jax supports something more generalized: their `in_dims` can match the structure of the `inputs` to the function (i.e., it is a nested python data structure matching the data structure of `inputs` specifying where in `inputs` the Tensors to be mapped are and what their map dims should be). We don't have the infrastruture yet so we only support `int` or a flat tuple for `in_dims`. Test Plan: - `pytest test/test_vmap.py -v` [ghstack-poisoned]

`in_dims` specifies which dimension of the input tensors should be vmapped over. One can also specify `None` as an `in_dim` for a particular input to indicate that we do not map over said input. We implement `in_dims` by creating a BatchedTensor with BatchDim equal to said `in_dim`. Most of this PR is error checking. `in_dims` must satisfy the following: - `in_dim` can be either an int or a Tuple[Optional[int]]. If it is an int, we use it to mean the `in_dim` for every input. - If `in_dims` is not-None at some index `idx`, then the input at index `idx` MUST be a tensor (vmap can only map over tensors). jax supports something more generalized: their `in_dims` can match the structure of the `inputs` to the function (i.e., it is a nested python data structure matching the data structure of `inputs` specifying where in `inputs` the Tensors to be mapped are and what their map dims should be). We don't have the infrastruture yet so we only support `int` or a flat tuple for `in_dims`. Test Plan: - `pytest test/test_vmap.py -v` ghstack-source-id: 36d06d9 Pull Request resolved: #40717

facebook-github-bot · 2020-07-07T04:05:35Z

@zou3519 merged this pull request in 5d1d8a5.

Summary: Pull Request resolved: pytorch#40717 `in_dims` specifies which dimension of the input tensors should be vmapped over. One can also specify `None` as an `in_dim` for a particular input to indicate that we do not map over said input. We implement `in_dims` by creating a BatchedTensor with BatchDim equal to said `in_dim`. Most of this PR is error checking. `in_dims` must satisfy the following: - `in_dim` can be either an int or a Tuple[Optional[int]]. If it is an int, we use it to mean the `in_dim` for every input. - If `in_dims` is not-None at some index `idx`, then the input at index `idx` MUST be a tensor (vmap can only map over tensors). jax supports something more generalized: their `in_dims` can match the structure of the `inputs` to the function (i.e., it is a nested python data structure matching the data structure of `inputs` specifying where in `inputs` the Tensors to be mapped are and what their map dims should be). We don't have the infrastruture yet so we only support `int` or a flat tuple for `in_dims`. Test Plan: - `pytest test/test_vmap.py -v` Differential Revision: D22397914 Pulled By: zou3519 fbshipit-source-id: 56d2e14be8b6024e4cde2729eff384da305b4ea3

This reverts commit 5ff9f58.

Summary: Pull Request resolved: pytorch#40717 `in_dims` specifies which dimension of the input tensors should be vmapped over. One can also specify `None` as an `in_dim` for a particular input to indicate that we do not map over said input. We implement `in_dims` by creating a BatchedTensor with BatchDim equal to said `in_dim`. Most of this PR is error checking. `in_dims` must satisfy the following: - `in_dim` can be either an int or a Tuple[Optional[int]]. If it is an int, we use it to mean the `in_dim` for every input. - If `in_dims` is not-None at some index `idx`, then the input at index `idx` MUST be a tensor (vmap can only map over tensors). jax supports something more generalized: their `in_dims` can match the structure of the `inputs` to the function (i.e., it is a nested python data structure matching the data structure of `inputs` specifying where in `inputs` the Tensors to be mapped are and what their map dims should be). We don't have the infrastruture yet so we only support `int` or a flat tuple for `in_dims`. Test Plan: - `pytest test/test_vmap.py -v` Differential Revision: D22397914 Pulled By: zou3519 fbshipit-source-id: 56d2e14be8b6024e4cde2729eff384da305b4ea3

This was referenced Jun 29, 2020

Initial vmap docstring #40575

Closed

Enable out_dims for vmap frontend API #40576

Closed

zou3519 requested review from cpuhrsch and ezyang June 29, 2020 21:19

ezyang reviewed Jul 6, 2020

View reviewed changes

ezyang approved these changes Jul 6, 2020

View reviewed changes

facebook-github-bot closed this in 5d1d8a5 Jul 7, 2020

facebook-github-bot added the merged label Jul 7, 2020

facebook-github-bot deleted the gh/zou3519/269/head branch July 10, 2020 14:18

csarofeen added a commit to csarofeen/pytorch that referenced this pull request Aug 16, 2020

Revert "Enable in_dims for vmap frontend api (pytorch#40717)"

19a6e31

This reverts commit 5ff9f58.

mruberry added the Merged label Oct 28, 2020

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Enable `in_dims` for vmap frontend api#40717

Enable `in_dims` for vmap frontend api#40717
zou3519 wants to merge 3 commits intogh/zou3519/269/basefrom
gh/zou3519/269/head

zou3519 commented Jun 29, 2020 •

edited

Loading

Uh oh!

dr-ci Bot commented Jun 29, 2020 •

edited

Loading

Uh oh!

ezyang Jul 6, 2020

Uh oh!

zou3519 Jul 6, 2020

Uh oh!

ezyang Jul 6, 2020

Uh oh!

zou3519 Jul 6, 2020

Uh oh!

ezyang Jul 6, 2020

Uh oh!

zou3519 Jul 6, 2020

Uh oh!

zou3519 Jul 6, 2020

Uh oh!

ezyang Jul 6, 2020

Uh oh!

zou3519 Jul 6, 2020

Uh oh!

facebook-github-bot commented Jul 7, 2020

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

zou3519 commented Jun 29, 2020 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

dr-ci Bot commented Jun 29, 2020 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

💊 CI failures summary and remediations

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot commented Jul 7, 2020

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

zou3519 commented Jun 29, 2020 •

edited

Loading

dr-ci Bot commented Jun 29, 2020 •

edited

Loading