torch.vmap is a prototype feature and should not be in the stable
binary. This PR:
- Removes the `torch.vmap` API
- Removes the documentation entry for `torch.vmap`
- Changes the vmap tests to use an internal API instead of `torch.vmap`.

Test Plan:
- Tested locally (test_torch, test_type_hints, test_vmap), but also wait
for CI.