torch.backends.nnpack 的源代码

from contextlib import contextmanager

import torch
from torch.backends import __allow_nonbracketed_mutation, ContextProp, PropModule

__all__ = ["is_available", "flags", "set_flags"]


def is_available():
    r"""返回 PyTorch 是否构建了 NNPACK 支持。"""
    return torch._nnpack_available()


def set_flags(_enabled):
    r"""设置是否全局启用 nnpack"""
    orig_flags = (torch._C._get_nnpack_enabled(),)
    torch._C._set_nnpack_enabled(_enabled)
    return orig_flags


@contextmanager
def flags(enabled=False):
    r"""用于设置是否全局启用 nnpack 的上下文管理器"""
    with __allow_nonbracketed_mutation():
        orig_flags = set_flags(enabled)
    try:
        yield
    finally:
        with __allow_nonbracketed_mutation():
            set_flags(orig_flags[0])