megenginelite-sys 1.8.2

A safe megenginelite wrapper in Rust
Documentation 
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
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55

import contextlib

from ..core._imperative_rt.core2 import (
    is_tracing_module,
    set_module_tracing,
    unset_module_tracing,
)

_enable_expr_checker = False
_enable_default_checker = True


def _get_expr_checker():
    return _enable_expr_checker


def _get_default_checker():
    return _enable_default_checker


def enable_expr_checker():
    r"""Call this function to check the result of each expr during tracing."""
    global _enable_expr_checker
    _enable_expr_checker = True
    _enable_default_checker = False


def disable_default_checker():
    r"""Call this function to disable checking the final output of the model after tracing."""
    global _enable_default_checker
    _enable_default_checker = False


_enable_graph_surgery_mode = False


def _graph_surgery_mode():
    return _enable_graph_surgery_mode


def _set_graph_surgery_mode(mode: bool):
    global _enable_graph_surgery_mode
    pre_mode = _enable_graph_surgery_mode
    _enable_graph_surgery_mode = mode
    return pre_mode


@contextlib.contextmanager
def _exclude_from_trace():
    is_tracing = is_tracing_module()
    if is_tracing:
        unset_module_tracing()
    yield
    if is_tracing:
        set_module_tracing()