[1.10] Remove torch.vmap (#65496)

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_autograd, test_type_hints, test_vmap),
but also wait for CI.

3 files changed, 2 insertions, 4 deletions

diff --git a/docs/source/torch.rst b/docs/source/torch.rst
index 084dad3490..c828f0ecbf 100644
--- a/docs/source/torch.rst
+++ b/docs/source/torch.rst
@@ -597,5 +597,4 @@ Utilities
     are_deterministic_algorithms_enabled
     set_warn_always
     is_warn_always_enabled
-    vmap
     _assert
diff --git a/test/test_vmap.py b/test/test_vmap.py
index 35b28db84f..6f6996b09b 100644
--- a/test/test_vmap.py
+++ b/test/test_vmap.py
@@ -1,7 +1,8 @@
 from torch.testing._internal.common_utils import TestCase, run_tests
 import torch
 import torch.nn.functional as F
-from torch import Tensor, vmap
+from torch import Tensor
+from torch._vmap_internals import vmap
 import functools
 import itertools
 import warnings
diff --git a/torch/__init__.py b/torch/__init__.py
index 18c9b08a8b..abd8c1f918 100644
--- a/torch/__init__.py
+++ b/torch/__init__.py
@@ -756,8 +756,6 @@ del register_after_fork
 # torch.jit.script as a decorator, for instance):
 from ._lobpcg import lobpcg as lobpcg
 
-from ._vmap_internals import vmap as 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 compatibility.