Computes the cache key given inputs and execution context.
(self, args, kwargs, include_tensor_ranks_only=False)
| 1950 | return self._descriptor_cache[instance] |
| 1951 | |
| 1952 | def _cache_key(self, args, kwargs, include_tensor_ranks_only=False): |
| 1953 | """Computes the cache key given inputs and execution context.""" |
| 1954 | if self.input_signature is None: |
| 1955 | inputs = (args, kwargs) if kwargs else args |
| 1956 | input_signature = pywrap_tensorflow.TFE_Py_EncodeArg( |
| 1957 | inputs, include_tensor_ranks_only) |
| 1958 | else: |
| 1959 | del args, kwargs |
| 1960 | assert not include_tensor_ranks_only |
| 1961 | input_signature = self.flat_input_signature |
| 1962 | |
| 1963 | ctx = context.context() |
| 1964 | |
| 1965 | # Don't need to open an init_scope if the _cache_key call is in eager mode |
| 1966 | # already. |
| 1967 | executing_eagerly = ctx.executing_eagerly() |
| 1968 | parent_graph = None |
| 1969 | if not executing_eagerly: |
| 1970 | with ops.init_scope(): |
| 1971 | # The graph, or whether we're executing eagerly, should be a part of the |
| 1972 | # cache key so we don't improperly capture tensors such as variables. |
| 1973 | executing_eagerly = ctx.executing_eagerly() |
| 1974 | parent_graph = None if executing_eagerly else ops.get_default_graph() |
| 1975 | |
| 1976 | # pylint: disable=protected-access |
| 1977 | default_graph = ops.get_default_graph() |
| 1978 | # TODO(b/117617952): The current distribution strategy will affect graph |
| 1979 | # building (e.g. accessing different variables from different devices) and |
| 1980 | # so requires retracing for each device. |
| 1981 | strategy_stack = default_graph._distribution_strategy_stack |
| 1982 | uses_distribution_strategy = ( |
| 1983 | strategy_stack and |
| 1984 | strategy_stack[-1].strategy.extended._retrace_functions_for_each_device |
| 1985 | ) |
| 1986 | if executing_eagerly: |
| 1987 | colocation_stack = () |
| 1988 | if uses_distribution_strategy: |
| 1989 | device_functions = (pydev.merge_device(ctx.device_name),) |
| 1990 | else: |
| 1991 | device_functions = () |
| 1992 | else: |
| 1993 | colocation_stack = tuple(default_graph._colocation_stack.peek_objs()) |
| 1994 | if (uses_distribution_strategy |
| 1995 | or func_graph_module.device_stack_has_callable( |
| 1996 | default_graph._device_function_stack)): |
| 1997 | # Putting the device in the cache key ensures that call-site device |
| 1998 | # annotations are respected. |
| 1999 | device_functions = tuple(default_graph._device_functions_outer_to_inner) |
| 2000 | else: |
| 2001 | device_functions = () |
| 2002 | |
| 2003 | in_cross_replica_context = False |
| 2004 | try: |
| 2005 | in_cross_replica_context = (strategy_stack[-1].replica_context is None) # pylint: disable=protected-access |
| 2006 | except (AttributeError, IndexError): |
| 2007 | pass |
| 2008 | |
| 2009 | return CacheKey(input_signature, parent_graph, device_functions, |
no test coverage detected