Context manager for copying distribute.Strategy scope information.
()
| 360 | |
| 361 | @tf_contextlib.contextmanager |
| 362 | def inner_cm(): |
| 363 | """Context manager for copying distribute.Strategy scope information.""" |
| 364 | graph = ops.get_default_graph() |
| 365 | # pylint: disable=protected-access |
| 366 | # TODO(b/112906995, nareshmodi): distribution strategy depends on |
| 367 | # inheriting this stack from the default graph even in eager mode. Maybe |
| 368 | # it should be part of the eager context? This would also allow us to |
| 369 | # remove a get_default_graph() call from the function cache lookup. |
| 370 | old_strategy_stack = self._distribution_strategy_stack |
| 371 | self._distribution_strategy_stack = list( |
| 372 | graph._distribution_strategy_stack) |
| 373 | # We ignore device placements from any outer scopes while tracing the |
| 374 | # function when possible, to avoid hard-coding them in the function |
| 375 | # graph. "Default" placements come from the PartitionedCallOp's placement, |
| 376 | # so that the same trace of the Python function may be placed on several |
| 377 | # different devices and saved functions may be placed on new devices when |
| 378 | # restored. |
| 379 | old_device_stack = self._device_function_stack |
| 380 | if context.executing_eagerly(): |
| 381 | if self._distribution_strategy_stack: |
| 382 | self._device_function_stack = self._device_function_stack.copy() |
| 383 | self._add_device_to_stack(context.context().device_name) |
| 384 | else: |
| 385 | if (self._distribution_strategy_stack |
| 386 | or device_stack_has_callable(graph._device_function_stack)): |
| 387 | # Hard-code devices from device functions in the function body |
| 388 | self._device_function_stack = graph._device_function_stack.copy() |
| 389 | |
| 390 | old_creator_stack = self._variable_creator_stack |
| 391 | self._variable_creator_stack = graph._variable_creator_stack |
| 392 | # Inherit the graph key, since this is used for matching variables in |
| 393 | # optimizers. |
| 394 | old_graph_key = self._graph_key |
| 395 | self._graph_key = graph._graph_key |
| 396 | # Inherit the auto_cast_variable_read_dtype, since this should not change |
| 397 | # inside a function. |
| 398 | old_auto_cast_var_read_dtype = self._auto_cast_variable_read_dtype |
| 399 | self._auto_cast_variable_read_dtype = graph._auto_cast_variable_read_dtype |
| 400 | # pylint: enable=protected-access |
| 401 | |
| 402 | with outer_cm as g: |
| 403 | try: |
| 404 | yield g |
| 405 | finally: |
| 406 | self._distribution_strategy_stack = old_strategy_stack |
| 407 | self._device_function_stack = old_device_stack |
| 408 | self._variable_creator_stack = old_creator_stack |
| 409 | self._graph_key = old_graph_key |
| 410 | self._auto_cast_variable_read_dtype = old_auto_cast_var_read_dtype |
| 411 | return inner_cm() |
| 412 | |
| 413 | @property |
nothing calls this directly
no test coverage detected