mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-14 20:57:59 +00:00
3ce1ebb6fb
Optimize unnecessary collection cast calls, unnecessary calls to list, tuple, and dict, and simplify calls to the sorted builtin. This should strictly improve speed and improve readability. Pull Request resolved: https://github.com/pytorch/pytorch/pull/94323 Approved by: https://github.com/albanD
26 lines
678 B
Python
26 lines
678 B
Python
# 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 = 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)
|