Add batching rule for Tensor.permute (#40517)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/40517
This is necessary for implementing the vmap frontend API's out_dims
functionality.
Test Plan:
- `./build/bin/vmap_test`. The vmap python API can't accept inputs that
aren't integers right now. There are workarounds around that (use a
lambda) but that doesn't look too nice. In the future we'll test all
batching rules in Python.
Differential Revision: D22216168
Pulled By: zou3519
fbshipit-source-id: b6ef552f116fddc433e242c1594059b9d2fe1ce4
