Batching rules for: chunk, split, unbind (#42480)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/42480
These are grouped together because they all return a tuple of multiple
tensors.
This PR implements batching rules for chunk, split, and unbind. It also
updates the testing logic. Previously, reference_vmap was not able to
handle multiple outputs, now, it does.
Test Plan: - `pytest test/test_vmap.py -v -k "Operators"`
Reviewed By: ezyang
Differential Revision: D22905401
Pulled By: zou3519
fbshipit-source-id: 9963c943d035e9035c866be74dbdf7ab1989f8c4
