Add BatchedTensorImpl (#38424)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/38424
On the way to adding initial vmap support, this is the implementation
for BatchedTensorImpl. Vmap (in future PRs) leverages Tensors backed by
BatchedTensorImpl to do its work.
For more context, here is an overview of the plan to add initial vmap support.
- [this PR] Add BatchedTensorImpl
- Add one or two batching rules
- Add vmap Python API
- Add "slow" for-loop fallbacks for out-of-place functions via
dispatcher fallback mechanism.
- Add batching rules for "view" functions
- Add "slow" for-loop fallbacks for in-place functions
- Miscellaneous handling for failure cases
- And more
Test Plan: - `./build/bin/vmap_test`
Differential Revision: D21640917
Pulled By: zou3519
fbshipit-source-id: 969490a838cf2099ed80104e7d51ee8ff069e168
