def testRewrapMutatesAffectedFunction(self): def new_target(x): return x * 3 self.assertEqual((1 * 2 + 1) ** 2, test_rewrappable_decorated(1)) prev_target, _ = tf_decorator.unwrap(test_rewrappable_decorated) tf_decorator.rewrap(test_rewrappable_decorated, prev_target, new_target) self.assertEqual((1 * 3 + 1) ** 2, test_rewrappable_decorated(1))
def testRewrapMutatesAffectedFunction(self): def new_target(x): return x * 3 self.assertEqual((1 * 2 + 1)**2, test_rewrappable_decorated(1)) prev_target, _ = tf_decorator.unwrap(test_rewrappable_decorated) tf_decorator.rewrap(test_rewrappable_decorated, prev_target, new_target) self.assertEqual((1 * 3 + 1)**2, test_rewrappable_decorated(1))
def testRewrapOfDecoratorFunction(self): def new_target(x): return x * 3 prev_target = test_rewrappable_decorated._tf_decorator._decorated_target # In this case, only the outer decorator (test_injectable_decorator_square) # should be preserved. tf_decorator.rewrap(test_rewrappable_decorated, prev_target, new_target) self.assertEqual((1 * 3) ** 2, test_rewrappable_decorated(1))
def testRewrapOfDecoratorFunction(self): def new_target(x): return x * 3 prev_target = test_rewrappable_decorated._tf_decorator._decorated_target # In this case, only the outer decorator (test_injectable_decorator_square) # should be preserved. tf_decorator.rewrap(test_rewrappable_decorated, prev_target, new_target) self.assertEqual((1 * 3)**2, test_rewrappable_decorated(1))
def tf_convert(f, ctx, convert_by_default=True, user_requested=False): """Decorator that applies AutoGraph to a function. Use in internal APIs. This API is suitable for high order functions internal to the TensorFlow API, and more generally any function to which Autograph is not applied. Guidance: convert was a decorator meant for use directly by developers, and will be soon deprecated in favor of tf.function. tf_convert is to be called from high order functions internal to TF. Args: f: Callable. ctx: ag_ctx.ControlStatusCtx, the Autograph context in which `f` is used. convert_by_default: bool, whether to use AutoGraph when the context doesn't specify. user_requested: bool, whether to ignore the conversion whitelist. See ConversionOptions.user_requested. Returns: Either `f or the converted version of `f`. """ if hasattr(f, '__ag_compiled'): return f f_wrapper = f decorators, f = tf_decorator.unwrap(f) # TODO(mdan): Grab features from context. # Note: we pass the original context through to convert to properly handle the # following scenario, which can be used insite TF implementations: # # ctx = ag_ctx.control_status_ctx() # @function(autograph=False) # Low-level graph code # def inner_fn(): # # The context is disabled here, but should be enabled in user user_fn # tf_convert(user_fn, ctx=ctx) if ctx.status == ag_ctx.Status.ENABLED: wrapper_factory = convert(recursive=True, user_requested=user_requested, conversion_ctx=ctx) elif ctx.status == ag_ctx.Status.DISABLED: wrapper_factory = do_not_convert elif ctx.status == ag_ctx.Status.UNSPECIFIED: if convert_by_default: wrapper_factory = convert(recursive=True, user_requested=user_requested, conversion_ctx=ctx) else: wrapper_factory = call_with_unspecified_conversion_status else: assert False, 'This switch contains all possible cases!' wrapper = wrapper_factory(f) if decorators: wrapper = tf_decorator.rewrap(f_wrapper, f, wrapper) setattr(wrapper, '__ag_compiled', True) return wrapper
def tf_convert(f, ctx, convert_by_default=True, force_conversion=False): """Decorator that applies AutoGraph to a function. Use in internal APIs. This API is suitable for high order functions internal to the TensorFlow API, and more generally any function to which Autograph is not applied. Guidance: convert was a decorator meant for use directly by developers, and will be soon deprecated in favor of tf.function. tf_convert is to be called from high order functions internal to TF. Args: f: Callable. ctx: ag_ctx.ControlStatusCtx, the Autograph context in which `f` is used. convert_by_default: bool, whether to use AutoGraph when the context doesn't specify. force_conversion: bool, whether to ignore the conversion whitelist. See ConversionOptions.force_conversion. Returns: Either `f or the converted version of `f`. """ if hasattr(f, '__ag_compiled'): return f f_wrapper = f decorators, f = tf_decorator.unwrap(f) # TODO(mdan): Grab features from context. if ctx.status == ag_ctx.Status.ENABLED: wrapper = convert(recursive=True, force_conversion=force_conversion)(f) elif ctx.status == ag_ctx.Status.DISABLED: wrapper = do_not_convert(f) elif ctx.status == ag_ctx.Status.UNSPECIFIED: if convert_by_default: wrapper = convert(recursive=True, force_conversion=force_conversion)(f) else: wrapper = call_with_unspecified_conversion_status(f) else: raise ValueError(ctx.status) if decorators: wrapper = tf_decorator.rewrap(f_wrapper, f, wrapper) setattr(wrapper, '__ag_compiled', True) return wrapper
def func_graph_from_py_func(name, python_func, args, kwargs, signature=None, func_graph=None, autograph=False, autograph_options=None, add_control_dependencies=True, arg_names=None, op_return_value=None, collections=None, capture_by_value=None): """Returns a `FuncGraph` generated from `python_func`. Args: name: an identifier for the function. python_func: the Python function to trace. args: the positional args with which the Python function should be called; ignored if a signature is provided. kwargs: the keyword args with which the Python function should be called; ignored if a signature is provided. signature: a possibly nested sequence of `TensorSpecs` specifying the shapes and dtypes of the arguments. When a signature is provided, `args` and `kwargs` are ignored, and `python_func` is traced with Tensors conforming to `signature`. If `None`, the shapes and dtypes are inferred from the inputs. func_graph: Optional. An instance of FuncGraph. If provided, we will use this graph else a new one is built and returned. autograph: whether to use autograph to compile `python_func`. See https://www.tensorflow.org/guide/autograph for more information. autograph_options: additional knobs to control when `autograph=True`. See https://www.tensorflow.org/guide/autograph for more information. add_control_dependencies: If True, automatically adds control dependencies to ensure program order matches execution order and stateful ops always execute. arg_names: Optional list of argument names, used to give input placeholders recognizable names. op_return_value: Optional. A Tensor. If set and `python_func` returns Operations, those return values will be replaced with this value. If not set, returning an Operation triggers an error. collections: a dictionary of collections this FuncGraph should start with. If not specified (None), the FuncGraph will read (but not write to) the outer graph's collections that are not whitelisted, and both read and write to the outer graph's collections that are whitelisted. The current whitelisted collections are the global variables, the local variables, and the trainable variables. Defaults to None. capture_by_value: An optional boolean. If True, the func graph will capture Variables by value instead of reference. By default inherit from outer graphs, and failing that will default to False. Returns: A FuncGraph. Raises: TypeError: If any of `python_func`'s return values is neither `None` nor a `Tensor`. """ if op_return_value is not None: assert isinstance(op_return_value, ops.Tensor), op_return_value if func_graph is None: func_graph = FuncGraph(name, collections=collections, capture_by_value=capture_by_value) assert isinstance(func_graph, FuncGraph) if add_control_dependencies: control_manager = AutomaticControlDependencies() else: control_manager = ops.NullContextmanager() with func_graph.as_default(), control_manager as a: current_scope = variable_scope.get_variable_scope() default_use_recource = current_scope.use_resource current_scope.set_use_resource(True) if signature is not None: args = signature kwargs = {} # Creates and names placeholders for all arguments. func_args = _get_defun_inputs_from_args(args, arg_names) func_kwargs = _get_defun_inputs_from_kwargs(kwargs) # Convert all Tensors into TensorSpecs before saving the structured inputs. # If storing pure concrete functions that are not called through polymorphic # functions, we don't have access to FunctionSpec, so we need to call the # TensorSpecs by their `arg_names` for later binding. func_graph.structured_input_signature = ( convert_structure_to_signature(func_args, arg_names), convert_structure_to_signature(func_kwargs)) flat_func_args = nest.flatten(func_args) flat_func_kwargs = nest.flatten(func_kwargs) # Temporarily set inputs to allow graph building code to inspect # them. Reassigned below. func_graph.inputs = [ arg for arg in flat_func_args + flat_func_kwargs if isinstance(arg, ops.Tensor) ] # Note: `nest.flatten` sorts by keys, as does `_deterministic_dict_values`. # Variables to help check whether mutation happens in calling the function # Copy the recursive list, tuple and map structure, but not base objects func_args_before = nest.pack_sequence_as(func_args, flat_func_args) func_kwargs_before = nest.pack_sequence_as(func_kwargs, flat_func_kwargs) def convert(x): """Converts a function output to a Tensor.""" if x is None: return None if op_return_value is not None and isinstance(x, ops.Operation): # TODO(b/79881896): we currently can't capture external control deps, so # this won't work if x needs to be captured (i.e. if python_func returns # captured Operations). with ops.control_dependencies([x]): x = array_ops.identity(op_return_value) elif not isinstance(x, tensor_array_ops.TensorArray): try: x = ops.convert_to_tensor_or_composite(x) except (ValueError, TypeError): raise TypeError( "To be compatible with tf.contrib.eager.defun, Python functions " "must return zero or more Tensors; in compilation of %s, found " "return value of type %s, which is not a Tensor." % (str(python_func), type(x))) if add_control_dependencies: x = a.mark_as_return(x) return x this_tape = tape.push_new_tape() try: if autograph: from tensorflow.python import autograph # pylint: disable=g-import-not-at-top _, original_func = tf_decorator.unwrap(python_func) def wrapper(*args, **kwargs): # Note: functions annotated with @tf.function should always be # converted even though they would meet autograph's whitelisting # criteria. # If this assumption is ever broken, converted_call will need to # handle the possibility of original_func still being a shim, e.g. # bound to WeakrefSelf. return autograph.converted_call( original_func, None, autograph.ConversionOptions( verbose=autograph.Verbosity.BRIEF, recursive=True, strip_decorators=(def_function.function, ), optional_features=autograph_options, force_conversion=True, ), args, kwargs) # Wrapping around a decorator allows checks like tf_inspect.getargspec # to be accurate. converted_func = tf_decorator.make_decorator( original_func, wrapper) tf_decorator.rewrap(python_func, original_func, converted_func) func_outputs = python_func(*func_args, **func_kwargs) # invariant: `func_outputs` contains only Tensors, IndexedSlices, # SparseTensors, TensorArrays and `None`s. func_outputs = nest.map_structure(convert, func_outputs) check_mutation(func_args_before, func_args) check_mutation(func_kwargs_before, func_kwargs) finally: tape.pop_tape(this_tape) current_scope.set_use_resource(default_use_recource) # Variables in `func_args`, `func_kwargs` should be explicit inputs # to the function, not captured inputs. tape_variables = this_tape.watched_variables() arg_variables = set() inputs = [] for arg in nest.flatten(func_args) + nest.flatten(func_kwargs): if isinstance(arg, resource_variable_ops.ResourceVariable): # Even if an argument variable was not used in the function, we've # already manually captured the resource Tensor when creating argument # placeholders. resource_placeholder = func_graph.captures.pop( arg.handle, None) if resource_placeholder is None: continue arg_variables.add(arg) inputs.append(resource_placeholder) elif isinstance(arg, ops.Tensor): inputs.append(arg) variables = [v for v in tape_variables if v not in arg_variables] func_graph.inputs = inputs + list(func_graph.captures.values()) func_graph.structured_outputs = func_outputs # Returning a closed-over tensor does not trigger convert_to_tensor. func_graph.outputs.extend( func_graph.capture(x) for x in flatten(func_graph.structured_outputs) if x is not None) func_graph.variables = variables if add_control_dependencies: func_graph.control_outputs.extend(control_manager.ops_which_must_run) # Register any other functions defined in the graph. with ops.init_scope(): if context.executing_eagerly(): for f in func_graph._functions.values(): # pylint: disable=protected-access # TODO(ashankar): What about the gradient registry? context.add_function(f._c_func.func) # pylint: disable=protected-access return func_graph
def func_graph_from_py_func(name, python_func, args, kwargs, signature=None, func_graph=None, autograph=False, add_control_dependencies=True, arg_names=None, op_return_value=None, collections=None): """Returns a `FuncGraph` generated from `python_func`. Args: name: an identifier for the function. python_func: the Python function to trace. args: the positional args with which the Python function should be called; ignored if a signature is provided. kwargs: the keyword args with which the Python function should be called; ignored if a signature is provided. signature: a possibly nested sequence of `TensorSpecs` specifying the shapes and dtypes of the arguments. When a signature is provided, `args` and `kwargs` are ignored, and `python_func` is traced with Tensors conforming to `signature`. If `None`, the shapes and dtypes are inferred from the inputs. func_graph: Optional. An instance of FuncGraph. If provided, we will use this graph else a new one is built and returned. autograph: whether to use autograph to compile `python_func`. See https://www.tensorflow.org/guide/autograph for more information. add_control_dependencies: If True, automatically adds control dependencies to ensure program order matches execution order and stateful ops always execute. arg_names: Optional list of argument names, used to give input placeholders recognizable names. op_return_value: Optional. A Tensor. If set and `python_func` returns Operations, those return values will be replaced with this value. If not set, returning an Operation triggers an error. collections: a dictionary of collections this FuncGraph should start with. If not specified (None), the FuncGraph will read (but not write to) the outer graph's collections that are not whitelisted, and both read and write to the outer graph's collections that are whitelisted. The current whitelisted collections are the global variables, the local variables, and the trainable variables. Defaults to None. Returns: A FuncGraph. Raises: TypeError: If any of `python_func`'s return values is neither `None` nor a `Tensor`. """ if op_return_value is not None: assert isinstance(op_return_value, ops.Tensor), op_return_value if func_graph is None: func_graph = FuncGraph(name, collections=collections) assert isinstance(func_graph, FuncGraph) if add_control_dependencies: control_manager = AutomaticControlDependencies else: control_manager = ops.NullContextmanager with func_graph.as_default(), control_manager() as a: current_scope = variable_scope.get_variable_scope() default_use_recource = current_scope.use_resource current_scope.set_use_resource(True) if signature is not None: args = signature kwargs = {} # Creates and names placeholders for all arguments. func_args = _get_defun_inputs_from_args(args, arg_names) func_kwargs = _get_defun_inputs_from_kwargs(kwargs) # Convert all Tensors into TensorSpecs before saving the structured inputs. # If storing pure concrete functions that are not called through polymorphic # functions, we don't have access to FunctionSpec, so we need to call the # TensorSpecs by their `arg_names` for later binding. func_graph.structured_input_signature = ( convert_structure_to_signature(func_args, arg_names), convert_structure_to_signature(func_kwargs)) # Note: `nest.flatten` sorts by keys, as does `_deterministic_dict_values`. # Variables to help check whether mutation happens in calling the function # Copy the recursive list, tuple and map structure, but not base objects func_args_before = nest.pack_sequence_as(func_args, nest.flatten(func_args)) func_kwargs_before = nest.pack_sequence_as( func_kwargs, nest.flatten(func_kwargs)) def convert(x): """Converts a function output to a Tensor.""" if x is None: return None if op_return_value is not None and isinstance(x, ops.Operation): # TODO(b/79881896): we currently can't capture external control deps, so # this won't work if x needs to be captured (i.e. if python_func returns # captured Operations). with ops.control_dependencies([x]): x = array_ops.identity(op_return_value) elif not isinstance(x, tensor_array_ops.TensorArray): try: x = ops.convert_to_tensor_or_composite(x) except (ValueError, TypeError): raise TypeError( "To be compatible with tf.contrib.eager.defun, Python functions " "must return zero or more Tensors; in compilation of %s, found " "return value of type %s, which is not a Tensor." % (str(python_func), type(x))) if add_control_dependencies: x = a.mark_as_return(x) return x this_tape = tape.push_new_tape() try: if autograph: from tensorflow.python import autograph # pylint: disable=g-import-not-at-top _, original_func = tf_decorator.unwrap(python_func) def wrapper(*args, **kwargs): # Note: functions annotated with @tf.function should always be # converted even though they would meet autograph's whitelisting # criteria. # If this assumption is ever broken, converted_call will need to # handle the possibility of original_func still being a shim, e.g. # bound to WeakrefSelf. return autograph.converted_call( original_func, None, autograph.ConversionOptions( verbose=autograph.Verbosity.BRIEF, recursive=True, strip_decorators=(def_function.function,), optional_features=(), force_conversion=True, ), *args, **kwargs) # Wrapping around a decorator allows checks like tf_inspect.getargspec # to be accurate. converted_func = tf_decorator.make_decorator(original_func, wrapper) tf_decorator.rewrap(python_func, original_func, converted_func) func_outputs = python_func(*func_args, **func_kwargs) # invariant: `func_outputs` contains only Tensors, IndexedSlices, # SparseTensors, TensorArrays and `None`s. func_outputs = nest.map_structure(convert, func_outputs) check_mutation(func_args_before, func_args) check_mutation(func_kwargs_before, func_kwargs) finally: tape.pop_tape(this_tape) current_scope.set_use_resource(default_use_recource) # Variables in `func_args`, `func_kwargs` should be explicit inputs # to the function, not captured inputs. tape_variables = this_tape.watched_variables() arg_variables = set() inputs = [] for arg in nest.flatten(func_args) + nest.flatten(func_kwargs): if isinstance(arg, resource_variable_ops.ResourceVariable): # Even if an argument variable was not used in the function, we've # already manually captured the resource Tensor when creating argument # placeholders. resource_placeholder = func_graph.captures.pop(arg.handle) arg_variables.add(arg) inputs.append(resource_placeholder) elif isinstance(arg, ops.Tensor): inputs.append(arg) variables = [v for v in tape_variables if v not in arg_variables] func_graph.inputs = inputs + list(func_graph.captures.values()) func_graph.structured_outputs = func_outputs # Returning a closed-over tensor does not trigger convert_to_tensor. func_graph.outputs.extend( func_graph.capture(x) for x in flatten(func_graph.structured_outputs) if x is not None) func_graph.variables = variables # Register any other functions defined in the graph. with ops.init_scope(): if context.executing_eagerly(): for f in func_graph._functions.values(): # pylint: disable=protected-access # TODO(ashankar): What about the gradient registry? context.add_function(f._c_func.func) # pylint: disable=protected-access return func_graph
def tf_convert(f, ctx, convert_by_default=True, user_requested=False): """Decorator that applies AutoGraph to a function. Use in internal APIs. This API is suitable for high order functions internal to the TensorFlow API, and more generally any function to which AutoGraph is not applied. Guidance: `convert` was a decorator meant for use directly by developers, but most of today's uses go through `tf.function`. `tf_convert` is to be called from high order functions internal to TF. By default, all the internal TensorFlow functions are skipped when AutoGraph processes the code. This may lead to user-supplied functions to be incorrectly skipped as well. `tf_convert` helps avoid that. See the following example for more details. ``` =====tf_internal_module.py===== def unconverted(input_fn): return input_fn() def converted(input_fn): return tf.__internal__.autograph.tf_convert( input_fn, ctx=tf.__internal__.autograph.control_status_ctx())() ======user_module.py====== @tf.function def foo(input_fn) return unconverted(input_fn) @tf.function def bar(input_fn) return converted(input_fn) @tf.function(autograph=False) def baz(input_fn) return converted(input_fn) ``` The `foo` method above will execute the `input_fn` without autograph conversion, while the `bar` method will run an autographed `input_fn`. The `baz` method will run an unconverted `input_fn`, since `tf_convert` respect the control status context. Note that both methods in `tf_internal_module` are skipped by autograph when tracing the `tf.function`. The configuration of whether a module/package should be skipped by autograph is controlled in tensorflow/python/autograph/core/config.py. Args: f: Callable. ctx: ag_ctx.ControlStatusCtx, the Autograph context in which `f` is used. convert_by_default: bool, whether to use AutoGraph when the context doesn't specify. user_requested: bool, whether to ignore the conversion allowlist. See ConversionOptions.user_requested. Returns: Either `f or the converted version of `f`. """ if is_autograph_artifact(f): return f f_wrapper = f decorators, f = tf_decorator.unwrap(f) # TODO(mdan): Grab features from context. # Note: we pass the original context through to convert to properly handle the # following scenario, which can be used inside TF implementations: # # ctx = ag_ctx.control_status_ctx() # @function(autograph=False) # Low-level graph code # def inner_fn(): # # The context is disabled here, but should be enabled in user user_fn # tf_convert(user_fn, ctx=ctx) if ctx.status == ag_ctx.Status.ENABLED: wrapper_factory = convert(recursive=True, user_requested=user_requested, conversion_ctx=ctx) elif ctx.status == ag_ctx.Status.DISABLED: wrapper_factory = do_not_convert elif ctx.status == ag_ctx.Status.UNSPECIFIED: if convert_by_default: wrapper_factory = convert(recursive=True, user_requested=user_requested, conversion_ctx=ctx) else: wrapper_factory = call_with_unspecified_conversion_status else: assert False, 'This switch contains all possible cases!' wrapper = wrapper_factory(f) if decorators: wrapper = tf_decorator.rewrap(f_wrapper, f, wrapper) return autograph_artifact(wrapper)
def func_graph_from_py_func(name, python_func, args, kwargs, signature=None, func_graph=None, experimental_autograph=False, add_control_dependencies=True, arg_names=None, op_return_value=None): """Returns a `FuncGraph` generated from `python_func`. Args: name: an identifier for the function. python_func: the Python function to trace. args: the positional args with which the Python function should be called; ignored if a signature is provided. kwargs: the keyword args with which the Python function should be called; ignored if a signature is provided. signature: a possibly nested sequence of `TensorSpecs` specifying the shapes and dtypes of the arguments. When a signature is provided, `args` and `kwargs` are ignored, and `python_func` is traced with Tensors conforming to `signature`. If `None`, the shapes and dtypes are inferred from the inputs. func_graph: Optional. An instance of FuncGraph. If provided, we will use this graph else a new one is built and returned. experimental_autograph: whether to use autograph to compile `python_func`. See https://www.tensorflow.org/guide/autograph for more information. add_control_dependencies: If True, automatically adds control dependencies to ensure program order matches execution order and stateful ops always execute. arg_names: Optional list of argument names, used to give input placeholders recognizable names. op_return_value: Optional. A Tensor. If set and `python_func` returns Operations, those return values will be replaced with this value. If not set, returning an Operation triggers an error. Returns: A FuncGraph. Raises: TypeError: If any of `python_func`'s return values is neither `None` nor a `Tensor`. """ if op_return_value is not None: assert isinstance(op_return_value, ops.Tensor), op_return_value if func_graph is None: func_graph = FuncGraph(name) assert isinstance(func_graph, FuncGraph) if add_control_dependencies: control_manager = AutomaticControlDependencies else: control_manager = ops.NullContextmanager with func_graph.as_default(), control_manager() as a: current_scope = variable_scope.get_variable_scope() default_use_recource = current_scope.use_resource current_scope.set_use_resource(True) if signature is not None: args = signature kwargs = {} # Creates and names placeholders for all arguments. func_args = _get_defun_inputs_from_args(args, arg_names) func_kwargs = _get_defun_inputs_from_kwargs(kwargs) # Note: `nest.flatten` sorts by keys, as does `_deterministic_dict_values`. # Variables to help check whether mutation happens in calling the function # Copy the recursive list, tuple and map structure, but not base objects func_args_before = nest.pack_sequence_as(func_args, nest.flatten(func_args)) func_kwargs_before = nest.pack_sequence_as( func_kwargs, nest.flatten(func_kwargs)) def convert(x): """Converts a function output to a Tensor.""" if x is None: return None if op_return_value is not None and isinstance(x, ops.Operation): # TODO(b/79881896): we currently can't capture external control deps, so # this won't work if x needs to be captured (i.e. if python_func returns # captured Operations). with ops.control_dependencies([x]): x = array_ops.identity(op_return_value) else: try: x = ops.convert_to_tensor_or_indexed_slices(x) except (ValueError, TypeError): raise TypeError( "To be compatible with tf.contrib.eager.defun, Python functions " "must return zero or more Tensors; in compilation of %s, found " "return value of type %s, which is not a Tensor." % (str(python_func), type(x))) if add_control_dependencies: x = a.mark_as_return(x) return x this_tape = tape.push_new_tape() try: if experimental_autograph: from tensorflow.python import autograph # pylint: disable=g-import-not-at-top _, original_func = tf_decorator.unwrap(python_func) def wrapper(*args, **kwargs): return autograph.converted_call( original_func, None, autograph.ConversionOptions( verbose=True, recursive=True, strip_decorators=(function.defun, def_function.function), optional_features=(), ), *args, **kwargs) # Wrapping around a decorator allows checks like tf_inspect.getargspec # to be accurate. converted_func = tf_decorator.make_decorator(original_func, wrapper) tf_decorator.rewrap(python_func, original_func, converted_func) func_outputs = python_func(*func_args, **func_kwargs) # invariant: `func_outputs` contains only Tensors and `None`s. func_outputs = nest.map_structure(convert, func_outputs) check_mutation(func_args_before, func_args) check_mutation(func_kwargs_before, func_kwargs) finally: tape.pop_tape(this_tape) current_scope.set_use_resource(default_use_recource) # Variables in `func_args`, `func_kwargs` should be explicit inputs # to the function, not captured inputs. tape_variables = this_tape.watched_variables() arg_variables = set() inputs = [] for arg in nest.flatten(func_args) + nest.flatten(func_kwargs): if isinstance(arg, resource_variable_ops.ResourceVariable): # Even if an argument variable was not used in the function, we've # already manually captured the resource Tensor when creating argument # placeholders. resource_placeholder = func_graph.captures.pop(arg.handle) arg_variables.add(arg) inputs.append(resource_placeholder) elif isinstance(arg, ops.Tensor): inputs.append(arg) variables = [v for v in tape_variables if v not in arg_variables] func_graph.inputs = inputs + list(func_graph.captures.values()) func_graph.structured_outputs = func_outputs # Returning a closed-over tensor does not trigger convert_to_tensor. func_graph.outputs.extend( func_graph.capture(x) for x in flatten(func_graph.structured_outputs) if x is not None) func_graph.variables = variables # Register any other functions defined in the graph. with ops.init_scope(): if context.executing_eagerly(): for f in func_graph._functions.values(): # pylint: disable=protected-access # TODO(ashankar): What about the gradient registry? context.add_function(f._c_func.func) # pylint: disable=protected-access return func_graph