"""Obsolete. Slated for deletion. Please use tf.function instead.

def __init__(self, *input_types, **kwargs):

"""Create a `Defun` decorator.

self._input_types = input_types

self._func_name = kwargs.pop("func_name", None)

self._grad_func = kwargs.pop("grad_func", None)

self._python_grad_func = kwargs.pop("python_grad_func", None)

self._out_names = kwargs.pop("out_names", None)

self._extra_kwargs = kwargs

def __call__(self, func):

# Various sanity checks on the callable func.

if not callable(func):

raise ValueError(f"Function {func} must be a callable.")

# Func should not use kwargs and defaults.

argspec = tf_inspect.getargspec(func)

if argspec.keywords or argspec.defaults:

raise ValueError(

"Functions with argument defaults or keywords arguments are not "

f"supported. {func} has defaults {argspec.defaults} and keywords "

f"{argspec.keywords}.")

# Computes how many arguments 'func' has.

min_args = len(argspec.args)

max_args = min_args

if argspec.varargs:

max_args = 1000000

argnames = argspec.args

if tf_inspect.ismethod(func):

# 1st argument is the "class" type.

min_args -= 1

argnames = argnames[1:]

if self._input_types:

# If Defun is given a list of types for the inputs, the number

# of input types should be compatible with 'func'.

num = len(self._input_types)

if num < min_args or num > max_args:

raise ValueError(

"The number of tf.function input types is not compatible with the "

f"allowed arguments of {func}. The tf.function have {num} input "

f"types, while the python function allows minimum {min_args} and "

f"maximum {max_args} arguments.")

return _DefinedFunction(

func,

argnames,

self._input_types,

self._func_name,

self._grad_func,

self._python_grad_func,

out_names=self._out_names,

**self._extra_kwargs)

# 'func' expects no arguments and input types is an empty list.

if min_args == 0 and max_args == 0:

return _DefinedFunction(

func, [], [],

self._func_name,

self._grad_func,

self._python_grad_func,

out_names=self._out_names,

**self._extra_kwargs)

# Input types are unknown. It's an overloaded function and hence

# its definition needs to be deferred until it's called.

return _OverloadedFunction(

func,

argnames,

self._func_name,

self._grad_func,

self._python_grad_func,

out_names=self._out_names,

"""Unregister function from eager context."""

__slots__ = ["name"]

def __init__(self, name):

self.name = name

def __del__(self):

try:

context.remove_function(self.name)

except TypeError:

# Suppress some exceptions, mainly for the case when we're running on

# module deletion. Things that can go wrong include the context module

# already being unloaded, self._handle._handle_data no longer being

# valid, and so on. Printing warnings in these cases is silly

# (exceptions raised from __del__ are printed as warnings to stderr).

pass # 'NoneType' object is not callable when the handle has been

# partially unloaded.

except AttributeError:

pass # 'NoneType' object has no attribute 'eager_mode' when context has

"""_DefinedFunction encapsulates a function definition and its properties.

def __init__(self,

func,

argnames,

input_types,

func_name=None,

grad_func=None,

python_grad_func=None,

out_names=None,

shape_func=None,

capture_by_value=False,

allowlisted_stateful_ops=None,

capture_resource_var_by_value=True,

**kwargs):

"""Creates _DefinedFunction.

self._func = func

self._input_types = input_types

self._func_name = func_name

self._grad_func = grad_func

self._python_grad_func = python_grad_func

self._out_names = out_names

self._shape_func = shape_func

self._capture_by_value = capture_by_value

self._allowlisted_stateful_ops = allowlisted_stateful_ops

if self._allowlisted_stateful_ops is None:

self._allowlisted_stateful_ops = set()

self._capture_resource_var_by_value = capture_resource_var_by_value

self._extra_kwargs = kwargs

# Constructed only when C API is disabled, lazily

self._definition = None

# Constructed only when C API is enabled, lazily

self._c_func = None

self._function_deleter = None

self._sub_functions = {} # Constructed with _definition or _c_func

# pylint: disable=protected-access

device_funcs = ops.get_default_graph()._device_functions_outer_to_inner

# pylint: enable=protected-access

# Get the innermost device if possible.

self._caller_device = device_funcs[-1] if device_funcs else None

# Cached OpDef for this function. When C API is enabled, this is

# the only part of FunctionDef that we cache in Python. When C API

# is disabled the whole _definition is available and this is simply

# another reference to _definition.signature

self._op_def = None

assert isinstance(input_types, (list, tuple))

self._arg_types = input_types

self._arg_names = [argnames[i] if i < len(argnames) else ("arg%d" % i)

for i in range(len(input_types))]

@property

def name(self):

"""Function name."""

self._create_definition_if_needed()

return self._func_name

@property

def cached_definition(self):

return self.definition

@property

def definition(self):

"""Function definition proto."""

self._create_definition_if_needed()

if self._c_func:

with c_api_util.tf_buffer() as buf:

with self._c_func.get() as func:

c_api.TF_FunctionToFunctionDef(func, buf)

fdef = function_pb2.FunctionDef()

proto_data = c_api.TF_GetBuffer(buf)

fdef.ParseFromString(compat.as_bytes(proto_data))

with ops.init_scope():

if context.executing_eagerly():

context.add_c_function(func)

self._function_deleter = _DefinedFunctionDeleter(

fdef.signature.name)

return fdef

return self._definition

@property

def _signature(self):

self._create_definition_if_needed()

return self._op_def

def set_grad_func(self, grad_func):

"""Specifies the gradient function of this function."""

assert not self._grad_func

assert isinstance(grad_func, _DefinedFunction)

self._grad_func = grad_func

@property

def grad_func_name(self):

"""Returns the name of the gradient function."""

return self._grad_func.name if self._grad_func else None

@property

def python_grad_func(self):

"""Python gradient function callable."""

return self._python_grad_func

@property

def declared_input_types(self):

"""Returns the list of data types of explicit declared inputs."""

return self._input_types

@property

def captured_inputs(self):

"""Returns the list of implicitly captured inputs."""

self._create_definition_if_needed()

return self._extra_inputs

@property

def stateful_ops(self):

"""Returns the list of stateful ops in function definition.

self._create_definition_if_needed()

return self._stateful_ops

def _create_definition_if_needed(self):

"""Creates the function definition if it's not created yet."""

with context.graph_mode():

self._create_definition_if_needed_impl()

def _create_definition_if_needed_impl(self):

"""This is not what you want, see _create_definition_if_needed."""

if self._definition is not None or self._c_func is not None:

return

# Copy variable collections (by reference) from the parent graph such that

# name based variable sharing (e.g. via tf.make_template) works between the

# func graph and parent graph.

variable_keys = []

variable_keys.extend(ops.GraphKeys._VARIABLE_COLLECTIONS) # pylint: disable=protected-access

variable_keys.append(vs._VARSTORE_KEY) # pylint: disable=protected-access

parent_graph = ops.get_default_graph()

collections_ref = {

key: parent_graph.get_collection_ref(key) for key in variable_keys}

temp_graph = func_graph_from_py_func(

self._func,

self._arg_names,

self._arg_types,

self._func_name,

self._capture_by_value,

self._caller_device,

collections_ref=collections_ref,

allowlisted_stateful_ops=self._allowlisted_stateful_ops,

capture_resource_var_by_value=self._capture_resource_var_by_value)

self._extra_inputs = temp_graph.extra_inputs

# pylint: disable=protected-access

self._sub_functions = temp_graph._functions

# pylint: enable=protected-access

# Extra kwargs are treated as attrs on the function def.

if self._func_name:

base_func_name = self._func_name

else:

base_func_name = function_utils.get_func_name(self._func)

if self._grad_func:

base_func_name += ("_%s" % self._grad_func.name)

kwargs_attr = _parse_kwargs_as_attrs(base_func_name, **self._extra_kwargs)

# FIXME(feyu): C API is always enabled now. The if-true branch never runs.

if not temp_graph._c_graph: # pylint: disable=protected-access

# Build the FunctionDef

self._definition = graph_to_function_def.graph_to_function_def(

temp_graph,

temp_graph.get_operations(),

temp_graph.inputs,

temp_graph.outputs,

out_names=self._out_names)

for k in kwargs_attr:

self._definition.attr[k].CopyFrom(kwargs_attr[k])

# Hash the definition and its dependencies.

self._hash_str = self._create_hash_str(

self._definition.signature.input_arg,

self._definition.signature.output_arg, self._definition.node_def)

# Finally, we decide the function name to use. If not specified,

# make up something which is almost certainly unique (but deterministic).

if not self._func_name:

self._func_name = "_".join([base_func_name, self._hash_str])

self._definition.signature.name = self._func_name

if self._func.__doc__:

self._definition.signature.description = self._func.__doc__

self._op_def = self._definition.signature

else: # C API is enabled

output_names = ([compat.as_bytes(x) for x in self._out_names]

if self._out_names else [])

description = self._func.__doc__ or None

# pylint: disable=protected-access

with temp_graph._c_graph.get() as c_graph:

c_func = c_api.TF_GraphToFunction_wrapper(

c_graph,

base_func_name,

self._func_name is None, # append_hash_to_fn_name

None, # opers

[t._as_tf_output() for t in temp_graph.inputs],

[t._as_tf_output() for t in temp_graph.outputs],

output_names,

[], # control_outputs

[], # control_output_names

None, # opts

description)

self._c_func = c_api_util.ScopedTFFunction(c_func, base_func_name)

# pylint: enable=protected-access

self._set_c_attrs(kwargs_attr)

# Set cached fields: _op_def and _func_name (if not already set)

self._op_def = self.definition.signature

if self._func_name:

assert self._func_name == self._op_def.name

else:

self._func_name = compat.as_str(self._op_def.name)

self._stateful_ops = [(op.name, op.type)

for op in temp_graph.get_operations()

if op._is_stateful] # pylint: disable=protected-access

def _set_c_attrs(self, attrs):

"""Sets `attrs` as attributes of self._c_func.

for name, attr_value in attrs.items():

serialized = attr_value.SerializeToString()

# TODO(skyewm): this creates and deletes a new TF_Status for every attr.

# It might be worth creating a convenient way to re-use the same status.

with self._c_func.get() as func:

c_api.TF_FunctionSetAttrValueProto(func, compat.as_str(name),

serialized)

def _create_hash_str(self, input_arg, output_arg, node_def):

"""Creates an 8-character string unique to this input.

hasher = hashlib.sha1()

def update_num(n):

hasher.update(compat.as_bytes("%x" % n))

def update_str(s):

update_num(len(s))

hasher.update(compat.as_bytes(s))

def update_strs(slist):

update_num(len(slist))

for s in slist:

update_str(s)

for adef in input_arg:

update_str(adef.SerializeToString())

for adef in output_arg:

update_str(adef.SerializeToString())

for n in sorted(node_def, key=lambda n: n.name):

update_str(n.name)

update_str(n.op)

update_strs(n.input)

update_num(len(n.attr))

# NOTE: protobuf map serialization does not guarantee ordering.

for k in sorted(n.attr):

update_str(k)

update_str(n.attr[k].SerializeToString())

return hasher.hexdigest()[:8]

def add_to_graph(self, g):

"""Adds this function into the graph g."""

self._create_definition_if_needed()

# Adds this function into 'g'.

# pylint: disable=protected-access

if context.executing_eagerly():

context.context().add_function_def(self.definition)

else:

g._add_function(self)

# pylint: enable=protected-access

# Ensures related sub-routines are defined in 'g', too.

for f in self._sub_functions.values():

g._add_function_recursive(f) # pylint: disable=protected-access

# Adds its gradient function, too.

if self._grad_func:

self._grad_func.add_to_graph(g)

def __call__(self, *args, **kwargs):

self.add_to_graph(ops.get_default_graph())

args = [ops.convert_to_tensor(_) for _ in args] + self._extra_inputs

ret, op = _call(self._signature, *args, **kwargs)

# Set a hidden attr in 'op' so that gradients_impl can refer back

# to this _DefinedFunction instance to access python_grad_func.

assert isinstance(op, ops.Operation)

setattr(op, "__defun", self)

if self._shape_func is not None:

shapes = self._shape_func(op)

if len(shapes) != len(op.outputs):

raise ValueError(f"shape_func {self._shape_func} produced "

f"{len(shapes):d} shapes, which does not match "

f"{len(op.outputs)} outputs.")

for (t, shape) in zip(op.outputs, shapes):

t.set_shape(shape)

"""_OverloadedFunction encapsulates an overloaded function.

def __init__(self,

func,

argnames,

func_name=None,

grad_func=None,

python_grad_func=None,

out_names=None,

**kwargs):

"""Creates _DefinedFunction.

self._func = func

self._argnames = argnames

self._func_name = func_name

assert grad_func is None or isinstance(grad_func, _OverloadedFunction)

self._grad_func = grad_func

self._python_grad_func = python_grad_func

self._out_names = out_names

self._extra_kwargs = kwargs

self._overload = {}

def instantiate(self, input_types):

"""Instantiate this function given input argument types.

# Stringify the type list.

key = _type_list_to_str(input_types)

defined = self._overload.get(key)

if not defined:

# If not defined yet, define the function given the input types.

name = self._func_name

if name is not None:

name = "_".join([name, key])

defined = _DefinedFunction(

self._func,

self._argnames,

input_types,

name,

None,

self._python_grad_func,

out_names=self._out_names,

**self._extra_kwargs)

_ = defined.name # Fully instantiate the function definition.

if self._grad_func:

# If _grad_func is given, it is another

# _OverloadedFunction. We need to instantiate it with the

# right input types.

output_types = [

dtypes.DType(_.type) for _ in defined._signature.output_arg # pylint: disable=protected-access

]

# pylint: disable=protected-access

defined._grad_func = self._grad_func.instantiate(input_types +

output_types)

# pylint: enable=protected-access

self._overload[key] = defined

return defined

def __call__(self, *args, **kwargs):

input_types = []

args = list(args)

for (i, x) in enumerate(args):

x = ops.convert_to_tensor(x)

if not isinstance(x, tensor_lib.Tensor):

raise ValueError(f"Expected a Tensor but got {x} with type {type(x)}.")

input_types.append(x.dtype)

args[i] = x

"""A helper for constructing a function.

def __init__(self, name, capture_by_value, allowlisted_stateful_ops,

capture_resource_var_by_value, *args, **kwargs):

super(_FuncGraph, self).__init__(*args, **kwargs)

self._capture_by_value = capture_by_value

self._allowlisted_stateful_ops = allowlisted_stateful_ops

self._capture_resource_var_by_value = capture_resource_var_by_value

self._building_function = True

self._outer_graph = ops.get_default_graph()

self._vscope = vs.get_variable_scope()

self._old_custom_getter = self._vscope.custom_getter

# The name of the function.

self.name = name

# Placeholder tensors representing the inputs to this function. The tensors

# are in this _FuncGraph.

self.inputs = []

# Tensors that will be returned this function. The tensors are in this

# _FuncGraph.

self.outputs = []

# Maps external tensor -> internal tensor (e.g. input placeholder).

self._captured = {}

# The external tensors that have been captured as inputs and must be passed

# to this function (empty if capturing by value, otherwise these are the

# keys of _captured).

self.extra_inputs = []

# Input placeholders that been added for captured values (empty if capturing

# by value).

self.extra_args = []

# Captured variables.

# TODO(skyewm): is this needed?

self.extra_vars = []

# pylint: disable=g-doc-return-or-yield

@property

def outer_graph(self):

"""The graph active when this _FuncGraph was created."""

return self._outer_graph

@tf_contextlib.contextmanager

def container(self, container_name):

"""Returns a context manager that specifies the resource container to use.

original_container = self._container

# pylint: disable=protected-access

with ops.init_scope():

original_init_container = ops.get_default_graph()._container

try:

self._container = container_name

with ops.init_scope():

ops.get_default_graph()._container = container_name

yield self._container

finally:

self._container = original_container

with ops.init_scope():

ops.get_default_graph()._container = original_init_container

# pylint: enable=protected-access

# pylint: enable=g-doc-return-or-yield

def getvar(

self,

getter,

name,

shape=None,

dtype=None,

initializer=None,

reuse=None,

trainable=True,

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

use_resource=None,

**kwargs):

"""A custom variable getter."""

# Here, we switch the default graph to the outer graph and ask the

# variable scope in which the function is defined to give us the

# variable. The variable is stashed in extra_vars and returned to

# the caller.

#

# We capture these variables so that the variable definition is

# hoisted upward to the outer most graph.

with self._outer_graph.as_default():

# pylint: disable=protected-access

var = self._vscope.get_variable(

vs._get_default_variable_store(),

name,

shape=shape,

dtype=dtype,

initializer=initializer,

reuse=reuse,

trainable=trainable,

collections=collections,

use_resource=use_resource)

self.extra_vars.append(var)

if (isinstance(var, resource_variable_ops.BaseResourceVariable) and

self._capture_resource_var_by_value):

# For resource-based variables read the variable outside the function

# and pass in the value. This ensures that the function is pure and

# differentiable. TODO(apassos) this may have performance problems if

# the function will only do embedding lookups on the variable.

return var.value()

return var

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):

for i, x in enumerate(inputs):

if isinstance(x, ops.EagerTensor) or x.graph is not self:

inputs[i] = self.capture(x)

return super(_FuncGraph, self)._create_op_internal(

op_type,

inputs,

dtypes=dtypes,

input_types=input_types,

name=name,

attrs=attrs,

op_def=op_def,

compute_device=compute_device)

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

"""Adds the given tensor to this graph and returns the captured tensor."""

if tensor.ref() in self._captured:

# Captured already.

return self._captured[tensor.ref()]

elif self._capture_by_value:

return self._add_tensor_and_parents(tensor)

else:

return self._capture_tensor_as_extra_input(tensor, name)

@property

def captures(self):

"""Pairs of tensors and captured tensor."""

return [(k.deref(), v) for k, v in self._captured.items()]

def _capture_tensor_as_extra_input(self, tensor, name=None):

# Substitute with a placeholder.

self.extra_inputs.append(tensor)

# Hoist the new input placeholder out of any control flow context

# we're currently in.

with ops.control_dependencies(None):

ph = array_ops.placeholder(

tensor.dtype, shape=tensor.get_shape(), name=name)

# pylint: disable=protected-access

if isinstance(tensor, ops.EagerTensor):

handle_data = tensor._handle_data

if handle_data:

handle_data = handle_data.SerializeToString()

else:

with tensor.graph._c_graph.get() as c_graph:

handle_data = c_api.GetHandleShapeAndType(c_graph,

tensor._as_tf_output())

if handle_data:

with ph.graph._c_graph.get() as c_graph:

c_api.SetHandleShapeAndType(c_graph, ph._as_tf_output(),

compat.as_bytes(handle_data))

# pylint: enable=protected-access

self.inputs.append(ph)

self._captured[tensor.ref()] = ph

self.extra_args.append(ph)

if _is_guaranteed_const(tensor):

with ops.control_dependencies(None):

return array_ops.guarantee_const(ph)

else:

return ph

def _add_tensor_and_parents(self, tensor):

op = self._add_op_and_parents(tensor.op)

return op.outputs[tensor.value_index]

def _add_op_and_parents(self, op: ops.Operation):

# pylint: disable=protected-access

op_def = graph_to_function_def._get_op_def(op)

if op._is_stateful and op not in self._allowlisted_stateful_ops:

raise ValueError(f"Cannot capture a stateful node (name:{op.name}, "

f"type:{op.type}) by value.")

elif op.type in ("Placeholder", "PlaceholderV2"):

raise ValueError(f"Cannot capture a placeholder (name:{op.name}, "

f"type:{op.type}) by value.")

# pylint: enable=protected-access

captured_inputs = [self._add_tensor_and_parents(x) for x in op.inputs]

captured_op = self._create_op_internal(

op.type,

captured_inputs, [o.dtype for o in op.outputs],

name=op.name,

attrs=op.node_def.attr,

op_def=op_def)

for t, captured_t in zip(op.outputs, captured_op.outputs):

self._captured[t.ref()] = captured_t

arg_names,

arg_types,

name=None,

capture_by_value=False,

device=None,

colocation_stack=None,

container=None,

collections_ref=None,

arg_shapes=None,

allowlisted_stateful_ops=None,

capture_resource_var_by_value=True):

"""Returns a _FuncGraph generated from `func`.

if not name:

name = function_utils.get_func_name(func)

func_graph = _FuncGraph(name, capture_by_value, allowlisted_stateful_ops,

capture_resource_var_by_value)

with func_graph.as_default(), ops.device(device):

# pylint: disable=protected-access

if collections_ref is not None:

func_graph._collections = collections_ref

if container is not None:

func_graph._container = container

if colocation_stack is not None:

func_graph._colocation_stack = colocation_stack

# pylint: enable=protected-access

if arg_shapes is None: