| 4980 | |
| 4981 | @property |
| 4982 | def _device_function_stack(self): |
| 4983 | if self._stack_state_is_thread_local: |
| 4984 | # This may be called from a thread where device_function_stack doesn't yet |
| 4985 | # exist. |
| 4986 | # pylint: disable=protected-access |
| 4987 | if not hasattr(self._thread_local, "_device_function_stack"): |
| 4988 | stack_copy_for_this_thread = self._graph_device_function_stack.copy() |
| 4989 | self._thread_local._device_function_stack = stack_copy_for_this_thread |
| 4990 | return self._thread_local._device_function_stack |
| 4991 | # pylint: enable=protected-access |
| 4992 | else: |
| 4993 | return self._graph_device_function_stack |
| 4994 | |
| 4995 | @property |
| 4996 | def _device_functions_outer_to_inner(self): |