From 0e857bf1092e0245925f6de596f7cb38de88cc23 Mon Sep 17 00:00:00 2001 From: Richard Zou Date: Fri, 24 Sep 2021 13:29:08 -0400 Subject: [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. --- docs/source/torch.rst | 1 - test/test_vmap.py | 3 ++- torch/__init__.py | 2 -- 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. -- cgit v1.2.3