File size: 668 Bytes
8193465 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 | # Copyright (c) Facebook, Inc. and its affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
from contextlib import contextmanager
from torch._C._functorch import _vmap_add_layers, _vmap_remove_layers
_enabled = False
@contextmanager
def _enable_layers(dims):
global _enabled
assert not _enabled
input = sorted((d._level, d.size) for d in dims if not isinstance(d, int))
n = len(input)
try:
_vmap_add_layers(input)
_enabled = True
yield
finally:
_enabled = False
_vmap_remove_layers(n)
|