[1.7] Remove torch.vmap (#45571)

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.

.jenkins/pytorch/common_utils.sh

@@ -52,9 +52,9 @@ function get_exit_code() {
 function file_diff_from_base() {
   # The fetch may fail on Docker hosts, but it's not always necessary.
   set +e
-  git fetch origin master --quiet
+  git fetch origin release/1.7 --quiet
   set -e
-  git diff --name-only "$(git merge-base origin/master HEAD)" > "$1"
+  git diff --name-only "$(git merge-base origin/release/1.7 HEAD)" > "$1"
 }
 
 function get_bazel() {

docs/source/torch.rst

@@ -535,5 +535,4 @@ Utilities
     promote_types
     set_deterministic
     is_deterministic
-    vmap
     Assert

test/test_vmap.py

@@ -1,6 +1,7 @@
 from torch.testing._internal.common_utils import TestCase, run_tests
 import torch
-from torch import Tensor, vmap
+from torch import Tensor
+from torch._vmap_internals import vmap
 import functools
 import warnings
 from torch.testing._internal.common_device_type import instantiate_device_type_tests

torch/__init__.py

@@ -608,8 +608,6 @@ del register_after_fork
 # torch.jit.script as a decorator, for instance):
 from ._lobpcg import lobpcg
 
-from ._vmap_internals import vmap
-
 # These were previously defined in native_functions.yaml and appeared on the
 # `torch` namespace, but we moved them to c10 dispatch to facilitate custom
 # class usage. We add these lines here to preserve backward compatbility.