Searched refs:colocation_stack (Results 1 – 3 of 3) sorted by relevance
34 colocation_stack: Any73 colocation_stack = ()79 colocation_stack = tuple(default_graph._colocation_stack.peek_objs())102 EagerContext(parent_graph, device_functions, colocation_stack,
943 colocation_stack=None, argument988 if colocation_stack is not None:989 func_graph._colocation_stack = colocation_stack
5420 def _colocation_stack(self, colocation_stack): argument5423 self._thread_local._colocation_stack = colocation_stack5426 self._graph_colocation_stack = colocation_stack