mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-25 11:46:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

make core.trace_state resetting be thread-local

Browse Source
This commit is contained in:
Matthew Johnson 2020-04-02 18:19:44 -07:00
parent b78b7a0309
commit 6d4987cc04
1 changed files with 1 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
jax/core.py
View File

@ -544,11 +544,10 @@ trace_state = TraceState()
def reset_trace_state() -> bool:
"Reset the global trace state and return True if it was already clean."
global trace_state
if (trace_state.substack != [Sublevel(0)] or
trace_state.trace_stack.downward or
trace_state.trace_stack.upward):
trace_state = TraceState()
trace_state.__init__()
return False
else:
return True