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
|
# 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 torch._C._functorch import (
_vmap_add_layers,
_vmap_remove_layers,
)
from contextlib import contextmanager
_enabled = False
@contextmanager
def _enable_layers(dims):
global _enabled
assert not _enabled
input = list(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)
|
