BatchedTensor fallback: extended to support ops with multiple Tensor returns (#42628)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/42628
This PR extends the BatchedTensor fallback to support operators with
multiple Tensor returns. If an operator has multiple returns, we stack
shards of each return to create the full outputs.
Test Plan:
- `pytest test/test_vmap.py -v`. Added a new test for an operator with
multiple returns (torch.var_mean).
Reviewed By: izdeby
Differential Revision: D22957095
Pulled By: zou3519
fbshipit-source-id: 5c0ec3bf51283cc4493b432bcfed1acf5509e662
