def start_subtrace():
"""Starts a nested trace, returns the Trace object."""
# TODO: This follows the __enter__ part of core.new_master. share
level = core.trace_state.trace_stack.next_level(False)
master = core.MasterTrace(level, pe.JaxprTrace)
core.trace_state.trace_stack.push(master, False)
return pe.JaxprTrace(master, core.cur_sublevel())
def start_subtrace(self):
"""Starts a nested trace, returns the Trace object."""
# TODO: This follows the __enter__ part of core.new_master.
if config.omnistaging_enabled:
level = core.thread_local_state.trace_state.trace_stack.next_level(
)
master = core.MasterTrace(level, pe.JaxprTrace)
core.thread_local_state.trace_state.trace_stack.push(master)
self._count_subtraces += 1
return pe.JaxprTrace(master, core.cur_sublevel())
else:
level = core.thread_local_state.trace_state.trace_stack.next_level(
False)
master = core.MasterTrace(level, pe.JaxprTrace)
core.thread_local_state.trace_state.trace_stack.push(master, False)
self._count_subtraces += 1
return pe.JaxprTrace(master, core.cur_sublevel())
