ops.GraphKeys.GLOBAL_VARIABLES,

ops.GraphKeys.LOCAL_VARIABLES,

ops.GraphKeys.TRAINABLE_VARIABLES,

variable_scope._VARSTORE_KEY, # pylint: disable=protected-access

variable_scope._VARSCOPESTORE_KEY # pylint: disable=protected-access

signature_context=None):

"""Convert a potentially nested structure to a signature.

def encode_arg(arg, path):

"""A representation for this argument, for converting into signatures."""

if isinstance(arg, tensor_lib.Tensor):

user_specified_name = None

try:

user_specified_name = compat.as_str(

arg.op.get_attr("_user_specified_name"))

except (ValueError, AttributeError):

pass

if path and user_specified_name and user_specified_name != path[0]:

# The user has explicitly named the argument differently than the name

# of the function argument.

name = user_specified_name

else:

name = tensor_lib.sanitize_spec_name("_".join(str(p) for p in path))

return tensor_lib.TensorSpec(arg.shape, arg.dtype, name)

if isinstance(arg, resource_variable_ops.ResourceVariable):

return trace_type.from_value(arg, signature_context)

if isinstance(arg, composite_tensor.CompositeTensor):

# TODO(b/133606651) Do we need to inject arg_name?

return arg._type_spec # pylint: disable=protected-access

if isinstance(arg, (

int,

float,

bool,

str,

type(None),

dtypes.DType,

tensor_lib.TensorSpec,

type_spec.TypeSpec,

)):

return arg

return UnknownArgument()

# We are using the flattened paths to name the TensorSpecs. We need an

# explicit name for them downstream.

flattened = nest.flatten_with_tuple_paths(structure)

if arg_names:

if len(arg_names) != len(structure):

raise ValueError(

"Passed in arg_names don't match actual signature (%s)." % arg_names)

# Replace all top-level names with their actual arg_names. If a path before

# was "(2,'a',1)", it will become "(arg_names[2],'a',1)".

flattened = [

((arg_names[path[0]],) + path[1:], arg) for path, arg in flattened

]

mapped = [encode_arg(arg, path) for path, arg in flattened]

"""Graph representing a function body.

def __init__(self,

name,

collections=None,

capture_by_value=None,

structured_input_signature=None,

structured_outputs=None):

"""Construct a new FuncGraph.

super().__init__()

self.name = name

# TODO(panzf): Separate captures from non-captures inputs in self.inputs

self.inputs = []

self.outputs = []

self.control_outputs = []

self.structured_input_signature = structured_input_signature

self.structured_outputs = structured_outputs

self._resource_tensor_inputs = object_identity.ObjectIdentitySet()

self._weak_variables = []

self._watched_variables = object_identity.ObjectIdentityWeakSet()

self.is_control_flow_graph = False

self._function_captures = capture_container.FunctionCaptures()

outer_graph = ops.get_default_graph()

self._weak_outer_graph = weakref.ref(outer_graph)

while outer_graph.building_function:

outer_graph = outer_graph.outer_graph

# If self._weak_outer_graph is deleted, we revert to the outermost Graph

# active when the FuncGraph was traced. This will not be a FuncGraph.

self._fallback_outer_graph = outer_graph

# If not None, records the names of output args of this function. Used to

# preserve the output names in the signature of a serialized+deserialized

# function. Private at the moment mostly because it's often out of date.

self._output_names = None

# Inherit capture-by-value from outer graph.

if capture_by_value is not None:

self.capture_by_value = capture_by_value

elif self.outer_graph is not None and isinstance(self.outer_graph,

FuncGraph):

self.capture_by_value = self.outer_graph.capture_by_value

else:

self.capture_by_value = False

self._building_function = True

graph = self.outer_graph

if context.executing_eagerly():

self.seed = context.global_seed()

# [for tf-data user migration from TF1.0 to 2.0] seed_used keep track of

# any None op_seed for random_op in the function, in which case we end up

# using function seed, which could be unintended behavior for the op.

self._seed_used = False

else:

self.seed = graph.seed

self._seed_used = False

# TODO(allenl): Figure out if we can remove colocation stack

# specialization (currently used in cond_v2), here and in the cache key.

self._colocation_stack = graph._colocation_stack.copy() # pylint: disable=protected-access

if collections is None:

for collection_name in graph.get_all_collection_keys():

if collection_name not in ALLOWLIST_COLLECTIONS:

self._collections[collection_name] = graph.get_collection(

collection_name)

for collection_name in ALLOWLIST_COLLECTIONS:

self._collections[collection_name] = graph.get_collection_ref(

collection_name)

else:

self._collections = collections

# Keep track of whether this FuncGraph is exportable to SavedModel. Use

# `graph.mark_as_unsaveable(reason)` to mark this FuncGraph and any

# dependent functions as unsaveable.

self._saveable = True

self._saving_errors = set()

# Keep track of callbacks to run when this graph exits default scope

self._scope_exit_callbacks = None

def __str__(self):

return "FuncGraph(name=%s, id=%s)" % (self.name, id(self))

def watch_variable(self, v):

"""Marks the variable v as accessed while building this graph."""

# Don't watch `v` if it is one of ResourceVariable input arguments.

if (isinstance(v, resource_variable_ops.ResourceVariable) and

v.handle in self._resource_tensor_inputs):

return

while self is not None and isinstance(self, FuncGraph):

self._watched_variables.add(v)

self = self.outer_graph

def capture_call_time_value(self,

closure,

spec,

key=None,

default_value=None,

placeholder=None):

"""Returns a placeholder which at call time has the value closure().

if key is None:

key = object()

if key not in self._function_captures.by_ref_internal:

trace_ctx = trace_type.InternalTracingContext(True)

spec = trace_type.from_value(spec, trace_ctx)

if placeholder is None:

placeholder_ctx = trace_type.InternalPlaceholderContext(self)

placeholder = spec.placeholder_value(placeholder_ctx)

def wrapped_closure():

# One major case requiring returning a `default_value` is when passing a

# concrete function to `save`, i.e.

# serving_fn = serve_fn.get_concrete_function(...)

# model.save(save_dir, signatures={"serving_default": serving_fn})

# `serving_fn` has deferred captures added through

# `capture_call_time_value`. It can't be saved correctly since

# `wrapped_closure` will end up executing under a default Graph instead

# of FuncGraph. The user of `capture_call_time_value` also cannot

# conditionally avoid this call since presence of `save_context` when

# executing `wrapped_closure` is not known at tracing time of

# `serving_fn`.

if save_context.in_save_context() and default_value is not None:

return default_value

# TODO(wxinyi): raise an error if in save context but no default value.

if not context.executing_eagerly():

graph = ops.get_default_graph()

assert isinstance(

graph, FuncGraph

), "This API should only be used in TF2 environment."

with graph.as_default():

ret_nest = graph.capture_call_time_value(

closure, spec, key=key, default_value=default_value)

else:

ret_nest = closure()

ret_nest = spec.cast(ret_nest, trace_type.InternalCastContext)

return spec.to_tensors(ret_nest)

wrapped_closure.output_spec = spec

self._function_captures.add_or_replace(

key=key,

external=wrapped_closure,

internal=placeholder,

tracetype=spec,

is_by_ref=True)

return self._function_captures.by_ref_internal[key]

def control_dependencies(self, control_inputs):

"""Handles control dependencies.

if control_inputs is None:

return super().control_dependencies(control_inputs)

filtered_control_inputs = []

for c in control_inputs:

# Check for _UnreadVariable

if (isinstance(c, indexed_slices.IndexedSlices) or

(hasattr(c, "_handle") and hasattr(c, "op"))):

c = c.op

graph_element = ops._as_graph_element(c) # pylint: disable=protected-access

if graph_element is None:

graph_element = c

if graph_element is not None and getattr(graph_element, "graph",

None) is not self:

self._function_captures.control.add(graph_element)

else:

filtered_control_inputs.append(graph_element)

return super().control_dependencies(filtered_control_inputs)

def as_default(self):

outer_cm = super().as_default()

return _func_graph_as_default_inner_cm(self, outer_cm)

@property

def outer_graph(self):

"""The Graph this FuncGraph is nested in.

current = self._weak_outer_graph()

if current is None:

return self._fallback_outer_graph

return current

@outer_graph.setter

def outer_graph(self, new_outer_graph):

"""Sets `outer_graph` to `new_outer_graph`."""

self._weak_outer_graph = weakref.ref(new_outer_graph)

@property

def output_types(self):

return [t.dtype for t in self.outputs]

@property

def output_shapes(self):

return [t.shape for t in self.outputs]

@property

def trainable_variables(self):

"""A sequence of trainable variables accessed by this FuncGraph.

return tuple(v for v in self.variables if v.trainable)

@property

def variables(self):

"""A sequence of variables accessed by this FuncGraph.

def deref(weak_v):

v = weak_v()

if v is None:

raise AssertionError(

"Called a function referencing variables which have been deleted. "

"This likely means that function-local variables were created and "

"not referenced elsewhere in the program. This is generally a "

"mistake; consider storing variables in an object attribute on "

"first call.")

return v

return tuple(deref(v) for v in self._weak_variables)

@variables.setter

def variables(self, var_list):

self._weak_variables = [weakref.ref(v) for v in var_list]

def _capture_by_value(

self,

op_type,

inputs,

dtypes, # pylint: disable=redefined-outer-name

input_types=None,

name=None,

attrs=None,

op_def=None,

compute_device=True):

# When capturing by value, do the read outside

reverse_captures = dict((id(v), k) for k, v in self.captures)

uncaptured_inputs = [reverse_captures.get(id(t), t) for t in inputs]

with ops.init_scope():

if context.executing_eagerly():

attr_list = ("dtype", int(attrs["dtype"].type))

value, = execute.execute(

compat.as_bytes(op_type), 1, uncaptured_inputs, attr_list,

context.context())

else:

op = ops.get_default_graph()._create_op_internal( # pylint: disable=protected-access

op_type, uncaptured_inputs, dtypes, input_types, name, attrs,

op_def, compute_device)

value = op.outputs[0]

captured_value = self.capture(value)

return captured_value.op

def _create_op_internal(

self,

op_type,

inputs,

dtypes=None, # pylint: disable=redefined-outer-name

input_types=None,

name=None,

attrs=None,

op_def=None,

compute_device=True):

"""Like Graph.create_op, except handles external input tensors.

if self.capture_by_value and op_type in [

"ReadVariableOp", "ResourceGather"

]:

return self._capture_by_value(op_type, inputs, dtypes, input_types, name,

attrs, op_def, compute_device)

# This capturing logic interacts poorly with control flow contexts which

# want to replace inputs of ops far too late in the process. This can lead

# the context to get confused and try to create an Enter for an Enter. We

# can detect this here and skip the additional Enter which can confuse loop

# validation logic.

if op_type == "Enter" and inputs[0].op.type == "Enter":

if inputs[0].op.get_attr("frame_name") == attrs["frame_name"].s:

return inputs[0].op

# Calling AddValue on the control flow contexts to force creation of the

# backward accumulators in the original graph before we create placeholders

# to capture the inputs.

ctxt = ops.get_default_graph()._control_flow_context # pylint: disable=protected-access

# Use a different list to avoid modifying the original inputs list.

captured_inputs = []

for inp in inputs:

# TPU Estimator defines a control flow context with no AddValue method.

if ctxt is not None and hasattr(ctxt, "AddValue"):

inp = ctxt.AddValue(inp)

inp = self.capture(inp)

captured_inputs.append(inp)

return super()._create_op_internal( # pylint: disable=protected-access

op_type, captured_inputs, dtypes, input_types, name, attrs, op_def,

compute_device)

def capture(self, tensor, name=None, shape=None):

return self._function_captures.capture_by_value(self, tensor, name)

def _validate_in_scope(self, tensor):

inner_graph = tensor.graph

while inner_graph is not None and isinstance(inner_graph, FuncGraph):

if inner_graph is self:

try:

tb = tensor.op.traceback

except AttributeError:

tensor_traceback = "<unknown>"

else:

tensor_traceback_list = []

for frame in traceback.format_list(tb.get_user_frames()):

tensor_traceback_list.extend(

[f" {line}" for line in frame.split("\n") if line.strip()])

tensor_traceback = "\n".join(tensor_traceback_list)

# Keep in sync with tfe_wrapper.cc.

# TODO(b/200991648): Unify those two paths.

raise errors.InaccessibleTensorError(

f"{tensor!r} is out of scope and cannot be used here. Use return "

"values, explicit Python locals or TensorFlow collections to "

"access it.\n"

"Please see https://www.tensorflow.org/guide/function#all_outputs_of_a_tffunction_must_be_return_values " # pylint: disable=line-too-long

"for more information.\n\n"

f"{tensor!r} was defined here:\n{tensor_traceback}\n\n"

f"The tensor {tensor!r} cannot be accessed from {self}, because "

f"it was defined in {tensor.graph}, which is out of scope.")

inner_graph = inner_graph.outer_graph

# TODO(panzf): Rename this method along with usages in cond/while graph.

def _capture_helper(self, tensor, name):

return self._function_captures._create_placeholder_helper( # pylint: disable=protected-access

self, tensor, name)

def _experimental_capture_side_input_by_ref(

self, identifier: Hashable, func: Callable[[], Any]

):

"""Implement capturing side input by reference for tf.function.

if context.executing_eagerly():

return func()

def maybe_convert_to_tensor():

value = func()

if not (isinstance(value, core.Value) or isinstance(value, core.Symbol)):

value = constant_op.constant(value)

return value

placeholder = self._function_captures._capture_by_ref( # pylint: disable=protected-access

self, maybe_convert_to_tensor, identifier)

return placeholder

@property

def captures(self):

"""Order list of tuples containing external and internal captures."""

return self._function_captures.by_val_capture_tuples

def add_capture(self, tensor, placeholder):

"""Capture a specific tensor and utilize the provided placeholder.

self._function_captures.add_or_replace(

key=id(tensor),

external=tensor,

internal=placeholder,

is_by_ref=False)

self.inputs.append(placeholder)

def replace_capture(self, tensor, placeholder):

"""Replace already existing capture."""

self._function_captures.add_or_replace(

key=id(tensor),

external=tensor,

internal=placeholder,

is_by_ref=False)

def replace_capture_with_deferred_capture(self,

tensor,

closure,

spec,

placeholder,

default_value=None):

"""Replaces existing capture `tensor` with a deferred capture `closure`.

self._function_captures.pop(id(tensor), is_by_ref=False)

self.capture_call_time_value(

closure,

spec,

key=id(tensor),

default_value=default_value,

placeholder=placeholder)

@property

def external_captures(self):

"""External tensors captured by this function."""

return list(self._function_captures.by_val_external.values())

@property

def internal_captures(self):

"""Placeholders in this function corresponding captured tensors."""

return list(self._function_captures.by_val_internal.values())

@property

def deferred_external_captures(self):

"""Ordered nest of tensors whose placeholders will be fed at call time."""

return list(self._function_captures.by_ref_external.values())

@property

def deferred_internal_captures(self):

"""List of nest of placeholders which at call time will be fed."""

return list(self._function_captures.by_ref_internal.values())

@property

def variable_captures(self):

"""Map of python object ids of variables to variables which are captured."""

return self.variables

@property

def function_captures(self):

return self._function_captures

def mark_as_unsaveable(self, error_message):

"""Marks this FuncGraph as unsaveable.

self._saveable = False

if isinstance(error_message, str):

error_message = [error_message]

self._saving_errors.update(error_message)

@property

def saveable(self):

"""Returns whether this FuncGraph is saveable."""

return self._saveable

@property

def saving_errors(self):

"""Returns set of errors preventing this FuncGraph from being saved."""

return self._saving_errors

def _add_scope_exit_callback(self, fn):

"""Add a function to call when this graph exits the default scope."""

if not callable(fn):

raise TypeError("fn is not callable: {}".format(fn))

if self._scope_exit_callbacks is None:

raise RuntimeError(

"Attempting to add a scope exit callback, but the default graph is "

"not the context scope graph. Did you forget to call "

"'with graph.as_default(): ...'?")

func_graph: FuncGraph, outer_cm: ContextManager[ops.Graph]):

"""Context manager for copying distribute.Strategy scope information."""

# pylint: disable=protected-access

# TODO(b/112906995, nareshmodi): distribution strategy depends on

# inheriting this stack from the default graph even in eager mode. Maybe

# it should be part of the eager context? This would also allow us to

# remove a get_default_graph() call from the function cache lookup.

graph = ops.get_default_graph()

old_strategy_stack = func_graph._distribution_strategy_stack

func_graph._distribution_strategy_stack = list(

graph._distribution_strategy_stack)

# We ignore device placements from any outer scopes while tracing the

# function when possible, to avoid hard-coding them in the function

# graph. "Default" placements come from the PartitionedCallOp's placement,

# so that the same trace of the Python function may be placed on several

# different devices and saved functions may be placed on new devices when

# restored.

# However, we need to preserve the outer device stack in the following

# cases in non eager context:

# 1. device stack is callable

# 2. When using distribution strategy with legacy graph mode.

old_device_stack = func_graph._device_function_stack

if (not context.executing_eagerly() and

(device_stack_has_callable(graph._device_function_stack) or

(func_graph._distribution_strategy_stack and

not ops.executing_eagerly_outside_functions()))):

# Hard-code devices from device functions in the function body

func_graph._device_function_stack = graph._device_function_stack.copy()

old_creator_stack = func_graph._variable_creator_stack

func_graph._variable_creator_stack = graph._variable_creator_stack

# Inherit the graph key, since this is used for matching variables in

# optimizers.

old_graph_key = func_graph._graph_key

func_graph._graph_key = graph._graph_key

old_scope_exit_callbacks = func_graph._scope_exit_callbacks

func_graph._scope_exit_callbacks = []

with outer_cm as g:

try:

yield g

finally:

try:

for fn in func_graph._scope_exit_callbacks:

fn()

finally:

func_graph._scope_exit_callbacks = old_scope_exit_callbacks

func_graph._distribution_strategy_stack = old_strategy_stack

func_graph._device_function_stack = old_device_stack

func_graph._variable_creator_stack = old_creator_stack

python_func,

args,

kwargs,

signature=None,

func_graph=None,

add_control_dependencies=True,

arg_names=None,

op_return_value=None,

collections=None,

capture_by_value=None,

create_placeholders=True):

"""Returns a `FuncGraph` generated from `python_func`.

if op_return_value is not None:

assert isinstance(op_return_value, tensor_lib.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:

deps_control_manager = auto_control_deps.AutomaticControlDependencies()

else:

deps_control_manager = ops.NullContextmanager()

with func_graph.as_default(), deps_control_manager as deps_ctx:

current_scope = variable_scope.get_variable_scope()

default_use_resource = current_scope.use_resource

current_scope.set_use_resource(True)

if signature is not None:

args = signature

kwargs = {}

if create_placeholders:

func_args, func_kwargs = _create_placeholders(args, kwargs, arg_names)

else:

func_args, func_kwargs = args, kwargs