"""Convert tensor t to an argdef, with a specified name or a unique name."""

arg = op_def_pb2.OpDef.ArgDef()

if name is None:

arg.name = _make_argname_from_tensor_name(t.name)

if used_names is not None:

if arg.name in used_names:

i = 0

while True:

new_name = "%s_U%d" % (arg.name, i)

if new_name not in used_names:

arg.name = new_name

break

i += 1

used_names.add(arg.name)

else:

arg.name = name

arg.type = t.dtype.as_datatype_enum

"""Create a mapping from graph tensor names to function tensor names."""

input_dict = {}

for op in function_graph.get_operations():

if _is_in_placeholders(op, func_arg_placeholders):

input_dict[op.values()[0].name] = op.values()[0].name

input_dict[op.name] = op.name

else:

op_def = _get_op_def(op)

attrs = _get_node_def(op).attr

o = 0

for arg_def in op_def.output_arg:

if arg_def.number_attr:

num = attrs[arg_def.number_attr].i

elif arg_def.type_list_attr:

num = len(attrs[arg_def.type_list_attr].list.type)

else:

num = 1

for i in range(num):

result = "%s:%s:%d" % (op.name, arg_def.name, i)

input_dict[op.values()[o].name] = result

if o == 0:

input_dict[op.name] = result

o += 1

"""Converts an op to a function def node and add it to `func`."""

# Add an entry in func.node_def

# Note that extend() makes a copy in this case, see:

# https://developers.google.com/protocol-buffers/docs/reference/python-generated#repeated-message-fields

func.node_def.extend([_get_node_def(op)])

node_def = func.node_def[-1]

for i in range(len(node_def.input)):

if not node_def.input[i].startswith("^"):

assert node_def.input[i] in input_dict, (

"%s missing from %s" % (node_def.input[i], input_dict.items()))

"""Returns `graph` as a `FunctionDef` protocol buffer.

func = function_pb2.FunctionDef()

func.signature.name = "_"

used_names = set()

func.signature.input_arg.extend([_tensor_to_argdef(i, used_names=used_names)

for i in inputs])

if out_names is None:

used_names = set()

func.signature.output_arg.extend([

_tensor_to_argdef(o, used_names=used_names) for o in outputs])

elif len(outputs) != len(out_names):

raise ValueError(

"Length of out_names (%d) does not match number of outputs (%d): %s" %

(len(out_names), len(outputs), ", ".join(out_names)))

elif len(out_names) != len(set(out_names)):

raise ValueError(

"Must not have duplicates in out_names: %s" % ", ".join(out_names))

else:

func.signature.output_arg.extend([

_tensor_to_argdef(o, name=n) for o, n in zip(outputs, out_names)])

func_arg_placeholders = set([i.name for i in inputs])

input_dict = _create_input_dict(graph, func_arg_placeholders)

for op in graph.get_operations():

if _is_in_placeholders(op, func_arg_placeholders):

continue

_add_op_node(op, func, input_dict)

if out_names is None:

for index, o in enumerate(outputs):

k = func.signature.output_arg[index].name

func.ret[k] = input_dict[o.name]

else:

for o, n in zip(outputs, out_names):

func.ret[n] = input_dict[o.name]

"""Parses **kwargs into a node's attributes."""

attrs = {}

noinline = kwargs.pop("noinline", None)

if noinline is not None:

attrs["_noinline"] = attr_value_pb2.AttrValue(b=bool(noinline))

compiled = kwargs.pop("compiled", None)

separate_compiled_gradients = kwargs.pop("separate_compiled_gradients", None)

if compiled is not None:

attrs["_XlaCompile"] = attr_value_pb2.AttrValue(b=bool(compiled))

attrs["_XlaSeparateCompiledGradients"] = attr_value_pb2.AttrValue(

b=bool(separate_compiled_gradients))

attrs["_XlaScope"] = attr_value_pb2.AttrValue(

s=("function_%s" % func_name).encode())

if kwargs:

raise ValueError("Unknown keyword arguments: %s" % kwargs.keys())

"""Adds a node calling a function.

if len(inputs) != len(sig.input_arg):

raise ValueError("Expected number of arguments: %d, received: %d" %

(len(sig.input_arg), len(inputs)))

name = kwargs.pop("name", None)

g = ops.get_default_graph()

func_name = sig.name

attrs = _parse_kwargs_as_attrs(func_name, **kwargs)

output_types = [dtypes.DType(x.type) for x in sig.output_arg]

with ops.name_scope(name, func_name, inputs) as name:

op = g.create_op(

func_name,

list(inputs),

output_types,

name=name,

attrs=attrs,

op_def=sig,

compute_shapes=False)

if op.outputs:

if len(op.outputs) == 1:

ret = op.outputs[0]

else:

ret = tuple(op.outputs)

else:

ret = op

if callable(func):

if inspect.isfunction(func):

return func.__name__

elif inspect.ismethod(func):

return "%s.%s" % (func.__self__.__name__, func.__name__)

else: # Probably a class instance with __call__

return type(func)

else:

"""A helper for construction a function.

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

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

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

self._captured = {}

self.extra_inputs = []

self.extra_args = []

self.extra_vars = []

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

# 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(self, op_type, inputs, data_types, **kwargs):

for i, x in enumerate(inputs):

if x.graph is not self:

# Referring to a tensor from other graph.

if x in self._captured:

# Captured already.

inputs[i] = self._captured[x]

else:

# Substitute with a placeholder.

self.extra_inputs.append(x)

ph = array_ops.placeholder(x.dtype, shape=x.get_shape())

# pylint: disable=protected-access

ph._handle_shape = x._handle_shape

ph._handle_dtype = x._handle_dtype

# pylint: enable=protected-access

inputs[i] = ph

self._captured[x] = ph

self.extra_args.append(ph)

return super(_FuncGraph, self).create_op(op_type, inputs, data_types,

"""Returns the captured variables by the function.

g = ops.get_default_graph()

if isinstance(g, _FuncGraph):

return g.extra_vars

else:

"""Returns the captured input tensors by the function.

g = ops.get_default_graph()

if isinstance(g, _FuncGraph):

return g.extra_inputs

else:

"""Returns the corresponding function arguments for the captured inputs.

g = ops.get_default_graph()

if isinstance(g, _FuncGraph):

return g.extra_args

else:

"""_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,

**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._extra_kwargs = kwargs

self._definition = None # Constructed lazily.

self._sub_functions = dict() # Constructed with definition.

self._args = []

assert isinstance(input_types, (list, tuple))

for i in range(len(input_types)):

argname = argnames[i] if i < len(argnames) else ("arg%d" % i)

argtype = input_types[i]

self._args.append((argname, argtype))

@property

def name(self):

"""Function name."""

self._create_definition_if_needed()

return self._func_name

@property

def definition(self):

"""Function definition proto."""

self._create_definition_if_needed()

return self._definition

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

"""Its gradient function's name."""

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

def _create_definition_if_needed(self):

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

if self._definition is not None:

return

# Create the func_def object.

temp_graph = _FuncGraph()

with temp_graph.as_default():

# List of placeholders for the function_def.

inputs = []

for (argname, argtype) in self._args:

argholder = array_ops.placeholder(argtype, name=argname)

inputs.append(argholder)

# Call func and gather the output tensors.

with vs.variable_scope("", custom_getter=temp_graph.getvar):

outputs = self._func(*inputs)

# If func only returned one value, make it a tuple.

if not isinstance(outputs, (list, tuple)):

outputs = (outputs,)

if any([_ is None for _ in outputs]):

raise ValueError("Function can not return None.")

# Ensures each output is a Tensor.

outputs = [ops.convert_to_tensor(_) for _ in outputs]

self._extra_inputs = temp_graph.extra_inputs

inputs.extend(temp_graph.extra_args)

# pylint: disable=protected-access

self._sub_functions = temp_graph._functions

# pylint: enable=protected-access

# Build the FunctionDef

self._definition = _graph_to_function_def(

temp_graph, inputs, outputs, out_names=self._out_names)

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

sig_pre_func_name = self._func_name or _get_func_name(self._func)

kwargs_attr = _parse_kwargs_as_attrs(

sig_pre_func_name, **self._extra_kwargs)

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([_get_func_name(self._func), self._hash_str])

self._definition.signature.name = self._func_name

if self._func.__doc__:

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

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

# pylint: disable=protected-access

# If 'g' has an identical function already, do nothing.

prev = g._get_function(self.name)

if prev and (prev._hash_str == self._hash_str):

return

# Adds this function into 'g'.

g._add_function(self)

# pylint: enable=protected-access

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

for f in self._sub_functions.values():

f.add_to_graph(g)

# 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._definition.signature, *args, **kwargs)

if self._shape_func is not None:

shapes = self._shape_func(op)

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

raise ValueError("shape_func produced %d shapes for %d outputs" %

(len(shapes), len(op.outputs)))

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

t.set_shape(shape)

"""Creates a _DefinedFunction initialized from a FunctionDef proto.

# The Python callable is only needed to create a FunctionDef. Since we have

# the FunctionDef here, we don't need to set _DefinedFunction._func (nor do we

# have access to such a callable here).

func = None

argnames = [arg.name for arg in fdef.signature.input_arg]

input_types = tuple(dtypes.as_dtype(arg.type)

for arg in fdef.signature.input_arg)

func_name = fdef.signature.name

# Note: FunctionDefs do not include python gradient functions, so if the

# original _DefinedFunction included one it will not be reflected here.

python_grad_func = None

out_names = [arg.name for arg in fdef.signature.output_arg]

result = _DefinedFunction(func, argnames, input_types, func_name, grad_func,

python_grad_func, out_names)

# pylint: disable=protected-access

result._definition = fdef

# Captured inputs are added as regular inputs to a function when it's

# serialized, i.e. any extra inputs from the original function are now

# included in `result`._args

result._extra_inputs = []

result._hash_str = result._create_hash_str(

result._definition.signature.input_arg,

result._definition.signature.output_arg,

result._definition.node_def)

# pylint: enable=protected-access

"""Creates _DefinedFunctions initialized from a FunctionDefLibrary proto.

if not lib.function and not lib.gradient: return []

# function name -> FunctionDef proto

funcs = {fdef.signature.name: fdef for fdef in lib.function}

# Validate that all references function names have function defs

for g in lib.gradient:

if g.function_name not in funcs:

raise ValueError("FunctionDefLibrary missing '%s' FunctionDef\n%s" %

(g.function_name, str(lib)))

if g.gradient_func not in funcs:

raise ValueError("FunctionDefLibrary missing '%s' FunctionDef\n%s" %

(g.gradient_func, str(lib)))

# function name -> gradient function name

func_to_grad = collections.defaultdict(lambda: None)

# gradient function name -> names of functions having that grad function

grad_to_funcs = collections.defaultdict(list)

for gdef in lib.gradient:

func_to_grad[gdef.function_name] = gdef.gradient_func

grad_to_funcs[gdef.gradient_func].append(gdef.function_name)

# Start with functions without gradients

ready = [fdef for fdef in lib.function

if func_to_grad[fdef.signature.name] is None]

if not ready:

raise ValueError("FunctionDefLibrary contains cyclic gradient functions!\n"

+ str(lib))

# function name -> _DefinedFunction

initialized = {}

while ready:

fdef = ready.pop()

name = fdef.signature.name

grad = initialized.get(func_to_grad[name])

if func_to_grad[name]: assert grad

defined_func = _from_definition(fdef, grad_func=grad)

initialized[name] = defined_func

ready.extend(funcs[f] for f in grad_to_funcs[name])

dtypes.float16: "f16",

dtypes.float32: "f32",

dtypes.float64: "f64",

dtypes.int32: "i32",

dtypes.uint8: "i8",

dtypes.uint16: "u16",

dtypes.int16: "i16",

dtypes.int8: "i8",

dtypes.string: "s",

dtypes.complex64: "c64",

dtypes.complex128: "c128",

dtypes.int64: "i64",

dtypes.bool: "b",

dtypes.qint8: "qi8",

dtypes.quint8: "qu8",

dtypes.qint16: "qi16",

dtypes.quint16: "qu16",

dtypes.qint32: "qi32",

dtypes.bfloat16: "b16"

"""_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.definition.signature.output_arg

]

# 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, ops.Tensor):

raise ValueError("Expect a Tensor but get ", x)

input_types.append(x.dtype)

args[i] = x

"""Decorator used to define TensorFlow functions.

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("func %s must be callable" % func)

# Func should not use kwargs and defaults.

argspec = inspect.getargspec(func)

if argspec.keywords or argspec.defaults:

raise ValueError("Functions with argument defaults or keyword "

"arguments are not supported.")

# 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 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 function has fewer arguments than the number of specified "

"input types.")

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,