Spaces:
Running
Running
File size: 3,260 Bytes
c61ccee |
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 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
import traceback
from contextlib import contextmanager
from typing import List, Any, Dict
from ._compatibility import compatibility
__all__ = ['preserve_node_meta', 'has_preserved_node_meta',
'set_stack_trace', 'set_grad_fn_seq_nr', 'reset_grad_fn_seq_nr',
'format_stack', 'set_current_meta', 'get_current_meta']
current_meta: Dict[str, Any] = {}
should_preserve_node_meta = False
@compatibility(is_backward_compatible=False)
@contextmanager
def preserve_node_meta():
global should_preserve_node_meta
saved_should_preserve_node_meta = should_preserve_node_meta
try:
should_preserve_node_meta = True
yield
finally:
should_preserve_node_meta = saved_should_preserve_node_meta
@compatibility(is_backward_compatible=False)
def set_stack_trace(stack : List[str]):
global current_meta
if should_preserve_node_meta and stack:
current_meta["stack_trace"] = "".join(stack)
@compatibility(is_backward_compatible=False)
def set_grad_fn_seq_nr(seq_nr):
global current_meta
if should_preserve_node_meta:
# The seq_nr is captured by eager mode in the grad_fn during forward
current_meta["grad_fn_seq_nr"] = current_meta.get("grad_fn_seq_nr", []) + [seq_nr]
current_meta["in_grad_fn"] = current_meta.get("in_grad_fn", 0) + 1
@compatibility(is_backward_compatible=False)
def reset_grad_fn_seq_nr():
# NB: reset state properly, this would be helpful towards supporting
# reentrant autograd if we actually wanted to do that.
global current_meta
if should_preserve_node_meta:
current_level = current_meta.get("in_grad_fn", 0)
assert current_level > 0
if current_level == 1:
del current_meta["in_grad_fn"]
del current_meta["grad_fn_seq_nr"]
else:
current_meta["in_grad_fn"] = current_level - 1
current_meta["grad_fn_seq_nr"].pop()
@compatibility(is_backward_compatible=False)
def format_stack() -> List[str]:
if should_preserve_node_meta:
return [current_meta.get("stack_trace", "")]
else:
# fallback to traceback.format_stack()
return traceback.format_list(traceback.extract_stack()[:-1])
@compatibility(is_backward_compatible=False)
def has_preserved_node_meta() -> bool:
return should_preserve_node_meta
@compatibility(is_backward_compatible=False)
@contextmanager
def set_current_meta(node):
global current_meta
if should_preserve_node_meta and node.meta:
saved_meta = current_meta
try:
current_meta = node.meta.copy()
# Append (node.name, node.target) onto "from_node" for provenance tracking
if "from_node" not in current_meta:
current_meta["from_node"] = [(node.name, node.target)]
elif current_meta["from_node"][-1][0] != node.name:
current_meta["from_node"].append((node.name, node.target))
yield
finally:
current_meta = saved_meta
else:
yield
@compatibility(is_backward_compatible=False)
def get_current_meta() -> Dict[str, Any]:
return current_meta
|