[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.

.circleci/scripts/binary_checkout.sh

@@ -61,7 +61,7 @@ git --no-pager log --max-count 1
 popd
 
 # Clone the Builder master repo
-retry git clone -q https://github.com/pytorch/builder.git "$BUILDER_ROOT"
+retry git clone -q https://github.com/pytorch/builder.git -b release/1.10 "$BUILDER_ROOT"
 pushd "$BUILDER_ROOT"
 echo "Using builder from "
 git --no-pager log --max-count 1

.jenkins/pytorch/common_utils.sh

@@ -99,5 +99,5 @@ function checkout_install_torchvision() {
 }
 
 function clone_pytorch_xla() {
-  git clone --recursive https://github.com/pytorch/xla.git
+  git clone --recursive -b r1.10 https://github.com/pytorch/xla.git
 }

docs/source/torch.rst

@@ -597,5 +597,4 @@ Utilities
     are_deterministic_algorithms_enabled
     set_warn_always
     is_warn_always_enabled
-    vmap
     _assert

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

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.