xref: /aosp_15_r20/external/pytorch/torch/_dynamo/current_scope_id.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1# mypy: allow-untyped-defs
2import contextlib
3import threading
4
5
6# Global variable to identify which SubgraphTracer we are in.
7# It is sometimes difficult to find an InstructionTranslator to use.
8_current_scope_id = threading.local()
9
10
11def current_scope_id():
12    global _current_scope_id
13    if not hasattr(_current_scope_id, "value"):
14        _current_scope_id.value = 1
15    return _current_scope_id.value
16
17
18@contextlib.contextmanager
19def enter_new_scope():
20    global _current_scope_id
21    try:
22        _current_scope_id.value = current_scope_id() + 1
23        yield
24    finally:
25        _current_scope_id.value = current_scope_id() - 1
26