# Copyright 2017 The TensorFlow Authors. All Rights Reserved.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

# ==============================================================================

# pylint: disable=unidiomatic-typecheck

"""Defun decorator for defining graph-mode functions."""

from __future__ import absolute_import

from __future__ import division

from __future__ import print_function

import collections

import functools

import re

import sys

import threading

import types as types_lib

import weakref

import numpy as np

import six

from tensorflow.core.framework import attr_value_pb2

from tensorflow.core.framework import function_pb2

from tensorflow.python import pywrap_tensorflow

from tensorflow.python.eager import context

from tensorflow.python.eager import execute

from tensorflow.python.eager import tape

from tensorflow.python.eager.graph_only_ops import graph_placeholder

from tensorflow.python.framework import c_api_util

from tensorflow.python.framework import constant_op

from tensorflow.python.framework import device as pydev

from tensorflow.python.framework import dtypes as dtypes_module

from tensorflow.python.framework import func_graph as func_graph_module

from tensorflow.python.framework import ops

from tensorflow.python.framework import tensor_spec

from tensorflow.python.ops import custom_gradient

from tensorflow.python.ops import functional_ops

from tensorflow.python.ops import gradients_impl

from tensorflow.python.ops import resource_variable_ops

from tensorflow.python.util import compat

from tensorflow.python.util import nest

from tensorflow.python.util import tf_decorator

from tensorflow.python.util import tf_inspect

# This is to avoid a circular dependency with gradients_impl

gradients_impl._function = sys.modules[__name__] # pylint: disable=protected-access

FORWARD_FUNCTION_ATTRIBUTE_NAME = "forward_function_name"

BACKWARD_FUNCTION_ATTRIBUTE_NAME = "backward_function_name"

# TODO(scottzhu): Update this to allow arbitrary attribute names in future.

WHITELIST_FUNCTION_ATTRIBUTE_REGEX = [

"experimental_.*",

FORWARD_FUNCTION_ATTRIBUTE_NAME,

BACKWARD_FUNCTION_ATTRIBUTE_NAME

]

def _parse_func_attrs(attributes):

"""Convert the keyword arguments into function_def attributes.

Currently only support primitive types: bool, int, float and string.

Args:

attributes: the dictionary of attributes.

Returns:

A dict of attributes where the key is the name of attribute and the value

is the AttrValue proto.

Raises:

ValueError: If the kwargs contains unwhitelisted name or unsupported value

types.

"""

attrs = {}

for key, value in attributes.items():

if not any([re.match(reg, key)

for reg in WHITELIST_FUNCTION_ATTRIBUTE_REGEX]):

raise ValueError("Attribute name is not whitelisted. "

"Whitelisted: prefix %s, got: %s" %

(WHITELIST_FUNCTION_ATTRIBUTE_REGEX, key))

if isinstance(value, attr_value_pb2.AttrValue):

attrs[key] = value

# bool type check has to happen before int since bool is a subclass of int.

elif isinstance(value, bool):

attrs[key] = attr_value_pb2.AttrValue(b=value)

elif isinstance(value, int):

attrs[key] = attr_value_pb2.AttrValue(i=value)

elif isinstance(value, float):

attrs[key] = attr_value_pb2.AttrValue(f=value)

elif isinstance(value, (str, bytes)):

attrs[key] = attr_value_pb2.AttrValue(s=compat.as_bytes(value))

else:

raise ValueError("Unsupported attribute type for %s with type %s" %

(key, type(value)))

return attrs

def _forward_name(n):

"""The name of a generated forward defun named n."""

return "__forward_%s_%s" % (n, ops.uid())

def _backward_name(n):

"""The name of a generated backward defun named n."""

return "__backward_%s_%s" % (n, ops.uid())

def _inference_name(n):

"""The name of a forward-but-no-gradient defun named n."""

return "__inference_%s_%s" % (n, ops.uid())

# TODO(apassos) get rid of this by splitting framework.function._DefinedFunction

# so it doesn't have the definition-generating logic and is just a container for

# an already-defined function.

class _EagerDefinedFunction(object):

"""Callable with the interface of `framework.function._DefinedFunction.`

`_EagerDefinedFunction` encapsulates a function definition and its properties,

and it provides a method for calling the encapsulated function. Some Ops

take functions as attributes, which have type `func`; an instance of this

class may be provided as the value of these `func` attributes.

"""

def __init__(self, name, graph, inputs, outputs, attrs):

"""Initializes an eager defined function.

Args:

name: str, the name for the created function.

graph: Graph, the graph containing the operations in the function

inputs: the tensors in the graph to be used as inputs to the function

outputs: the tensors in the graph which will be outputs to the function

attrs: dict mapping names of attributes to their AttrValue values

"""

operations = [

op for op in graph.get_operations()

if op not in set(arg.op for arg in inputs)

]

fn = pywrap_tensorflow.TF_GraphToFunction_wrapper(

graph._c_graph, # pylint: disable=protected-access

compat.as_str(name),

False,

[o._c_op for o in operations], # pylint: disable=protected-access

[t._as_tf_output() for t in inputs], # pylint: disable=protected-access

[t._as_tf_output() for t in outputs], # pylint: disable=protected-access

[],

None,

compat.as_str(""))

for name, attr_value in attrs.items():

serialized = attr_value.SerializeToString()

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

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

pywrap_tensorflow.TF_FunctionSetAttrValueProto(

fn, compat.as_str(name), serialized)

# TODO(apassos) avoid creating a FunctionDef (specially to grab the

# signature, but also in general it's nice not to depend on it.

with c_api_util.tf_buffer() as buffer_:

pywrap_tensorflow.TF_FunctionToFunctionDef(fn, buffer_)

proto_data = pywrap_tensorflow.TF_GetBuffer(buffer_)

function_def = function_pb2.FunctionDef()

function_def.ParseFromString(compat.as_bytes(proto_data))

with ops.init_scope():

if context.executing_eagerly():

context.add_function(fn)

self.definition = function_def

self.name = compat.as_bytes(function_def.signature.name)

self.signature = function_def.signature

self._num_outputs = len(self.signature.output_arg)

self._output_types = [o.type for o in self.signature.output_arg]

self._output_shapes = [o.shape for o in outputs]

self._func_graph_outputs = outputs

self.grad_func_name = None

self.python_grad_func = None

self._c_func = c_api_util.ScopedTFFunction(fn)

self._grad_func = None

self._graph = graph

self._stateful_ops = tuple(op for op in operations if op.op_def.is_stateful)

def add_to_graph(self, g):

# pylint: disable=protected-access

if self.name not in g._functions:

g._add_function(self)

for f in self._graph._functions.values():

if f.name not in g._functions:

g._add_function(f)

# pylint: enable=protected-access

@property

def stateful_ops(self):

return self._stateful_ops

def call(self, ctx, args):

"""Calls this function with `args` as inputs.

Function execution respects device annotations only if the function won't

be compiled with xla.

Args:

ctx: a Context object

args: a list of arguments to supply this function with.

Returns:

The outputs of the function call.

Raises:

ValueError: if the number of arguments is incorrect.

"""

executing_eagerly = ctx.executing_eagerly()

if self._graph._xla_compile: # pylint: disable=protected-access

# XLA compilation relies upon a custom kernel creator to run functions.

signature = self.signature

if executing_eagerly:

outputs = execute.execute(

str(signature.name),

num_outputs=self._num_outputs,

inputs=args,

attrs=None,

ctx=ctx)

else:

g = ops.get_default_graph()

self.add_to_graph(g)

op = g.create_op(

signature.name,

[ops.internal_convert_to_tensor(x, ctx=ctx) for x in args],

tuple(dtypes_module.DType(x.type) for x in signature.output_arg),

op_def=signature,

name="FunctionCall",

compute_shapes=False)

outputs = op.outputs

if not outputs:

return op

outputs = [outputs] if isinstance(

outputs, (ops.Tensor, type(None))) else list(outputs)

else:

# TODO(akshayka): Either remove this if the FunctionLibraryRuntime

# creates `PartitionedCallOp` kernels by default, or remove the previous

# branch if a TPU kernel is registered for `PartitionedCall`.

if len(args) != len(self.signature.input_arg):

raise ValueError(

"Arguments and signature arguments do not match: %s %s " %

(len(args), len(list(self.signature.input_arg))))

outputs = functional_ops.partitioned_call(

args=args,

f=self,

tout=self._output_types,

executing_eagerly=executing_eagerly,

config=ctx.rewriter_config_string) # pylint: disable=protected-access

if executing_eagerly:

return outputs

else:

for i, shape in enumerate(self._output_shapes):

outputs[i].set_shape(shape)

for i, func_graph_output in enumerate(self._func_graph_outputs):

custom_gradient.copy_handle_data(func_graph_output, outputs[i])

return outputs

class Function(object):

"""Callable object encapsulating a function definition and its gradient.

`Function` is a callable that encapsulates a function definition and

is differentiable under `tf.GradientTape` objects.

"""

def __init__(self, func_graph, attrs=None, signature=None):

"""Initialize a Function.

Args:

func_graph: An instance of FuncGraph: the function body to wrap.

attrs: (optional) dict mapping names of attributes to their AttrValue

values. Attributes in `attrs` will be included in this function's

definition.

Raises:

ValueError: If number of input_placeholders is not equal to the number

of function inputs.

"""

self._func_graph = func_graph

self._captured_inputs = list(self._func_graph.captures.keys())

self._num_outputs = len(self._func_graph.outputs)

self._output_shapes = tuple(

output.shape for output in self._func_graph.outputs)

self._attrs = _parse_func_attrs(attrs or {})

self._inference_function = _EagerDefinedFunction(

_inference_name(self._func_graph.name), self._func_graph,

self._func_graph.inputs, self._func_graph.outputs, self._attrs)

self._backward_graph_function = None

self._signature = signature

self._gradient_name = None

def __call__(self, *args):

"""Executes the wrapped function.

Args:

*args: a list of Tensors or Variables.

Returns:

The result of applying the TF function to `args`.

Raises:

ValueError: If `args` contains anything other than Tensors or Variables.

"""

ctx = context.context()

for v in self._func_graph.variables:

if v.trainable:

tape.variable_accessed(v)

tensor_inputs = []

for i, arg in enumerate(nest.flatten(args)):

if isinstance(arg, resource_variable_ops.ResourceVariable):

if arg.trainable:

tape.variable_accessed(arg)

tensor_inputs.append(arg.handle)

elif isinstance(arg, ops.Tensor):

tensor_inputs.append(arg)

elif (self._signature is not None and

isinstance(self._signature[i], tensor_spec.TensorSpec)):

tensor_inputs.append(

ops.convert_to_tensor(arg, self._signature[i].dtype))

else:

raise ValueError("All inputs to `Function`s must be Tensors; "

"on invocation of %s, the %d-th input (%s) was not a "

"Tensor." % (self._func_graph.name, i, str(arg)))

args = tensor_inputs + self._captured_inputs

if (tape.should_record(tensor_inputs) or

tape.should_record(self._captured_inputs)):

return self._backprop_call(args)

# Only need to override the gradient in graph mode and when we have outputs.

if context.executing_eagerly() or not self.outputs:

outputs = self._inference_function.call(ctx, args)

else:

if not self._gradient_name:

self._gradient_name = "PartitionedCall-%s" % ops.uid()

self._register_gradient(self._gradient_name)

with ops.get_default_graph().gradient_override_map(

{"PartitionedCall": self._gradient_name,

"StatefulPartitionedCall": self._gradient_name}):

outputs = self._inference_function.call(ctx, args)

return self._build_call_outputs(outputs)

def _register_gradient(self, name):

"""Registers the gradient for the current Function under the given name.

The gradient rewrites an inference call op to a forward call op, but does

not modify a pre-existing forward call op. It then computes the gradient

from the output's gradients and the side outputs of the forward op.

Args:

name: The name to register the gradient as.

"""

@ops.RegisterGradient(name)

def grad_fn(op, *doutputs): # pylint: disable=unused-variable

"""Gradients of this function."""

if self._backward_graph_function is None:

self._construct_backprop_function()

# pylint: disable=protected-access

self._forward_function.add_to_graph(op.graph)

num_inference_outputs = self._inference_function._num_outputs

# Rewrite an inference call op to be a forward call op

if op.get_attr("f").name.encode() == self._inference_function.name:

func = attr_value_pb2.AttrValue(

func=attr_value_pb2.NameAttrList(

name=self._forward_function.name))

op._set_attr("f", func)

types = attr_value_pb2.AttrValue.ListValue(

type=self._forward_function._output_types)

op._set_attr("Tout", attr_value_pb2.AttrValue(list=types))

for i in range(

num_inference_outputs, len(self._forward_function._output_types)):

t = ops.Tensor(op, i, self._forward_function._output_types[i])

t.set_shape(self._forward_function._output_shapes[i])

func_graph_output = self._forward_function._func_graph_outputs[i]

custom_gradient.copy_handle_data(func_graph_output, t)

op._outputs.append(t)

# pylint: enable=protected-access

# Compute the gradients using the side outputs

side_outputs = op.outputs[num_inference_outputs:]

args = list(doutputs[:num_inference_outputs]) + list(side_outputs)

return self._backward_graph_function(*[a for a in args if a is not None])

@property

def name(self):

"""Function name."""

return self._inference_function.name

@property

def graph(self):

"""Returns the graph from which this function was constructed."""

return self._func_graph

@property

def inputs(self):

"""Returns tensors in `self.graph` corresponding to arguments."""

return self._func_graph.inputs

@property

def outputs(self):

"""Returns tensors in `self.graph` corresponding to return values."""

return self._func_graph.outputs

@property

def captured_inputs(self):

"""Returns external Tensors captured by this function.

self.__call__(*args) passes `args + self.captured_inputs` to the function.

"""

return self._captured_inputs

@property

def function_def(self):

"""Returns a `FunctionDef` object representing this function."""

return self._inference_function.definition

@property

def output_shapes(self):

"""The function's output shapes."""

# TODO(ebrevdo): Should we only keep the output shapes associated

# with len(self._python_returns) outputs?

# TODO(akshayka): Consider removing this.

outputs_list = nest.flatten(self._func_graph.structured_outputs)

j = 0

for i, o in enumerate(outputs_list):

if o is not None:

if isinstance(o, ops.IndexedSlices):

# Extract the shape of the `IndexedSlices` object's `values` field.

outputs_list[i] = self._output_shapes[j] # the `values` shape

if o.dense_shape is not None:

j += 3 # skip over shapes for `values`, `indices`, `dense_shape`

else:

j += 2 # skip over shapes for `values`, `indices`

else:

outputs_list[i] = self._output_shapes[j]

j += 1

return nest.pack_sequence_as(self._func_graph.structured_outputs,

outputs_list)

@property

def output_dtypes(self):

# TODO(akshayka): Consider removing this.

return nest.map_structure(lambda x: x.dtype if x is not None else None,

self._func_graph.structured_outputs)

def add_to_graph(self, g):

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

return self._inference_function.add_to_graph(g)

def _construct_backprop_function(self):

"""Constructs the backprop function object for this function."""

backwards_graph = func_graph_module.FuncGraph(

_backward_name(self._func_graph.name))

forward_function_name = _forward_name(self._func_graph.name)

outputs = [x for x in self._func_graph.outputs

if gradients_impl.IsTrainable(x)]

with backwards_graph.as_default():

gradients_wrt_outputs = [

graph_placeholder(x.dtype, x.shape) for x in outputs

]

gradients_wrt_inputs = gradients_impl._GradientsHelper( # pylint: disable=protected-access

outputs,

self._func_graph.inputs,

grad_ys=gradients_wrt_outputs,

src_graph=self._func_graph,

unconnected_gradients=gradients_impl.UnconnectedGradients.NONE)

backwards_graph_captures = list(backwards_graph.captures.keys())

backward_function_attr = _parse_func_attrs(

{FORWARD_FUNCTION_ATTRIBUTE_NAME: forward_function_name})

backward_function_attr.update(self._attrs)

# The ordering of `backwards_graph.inputs` is important: inputs of

# `self._backward_graph_function` correspond to outputs of

# `self._forward_function`.

backwards_graph.inputs = gradients_wrt_outputs + list(

backwards_graph.captures.values())

# Clear captures, since we pass them in as inputs.

backwards_graph.captures = {}

backwards_graph.outputs.extend(

grad for grad in func_graph_module.flatten(gradients_wrt_inputs)

if grad is not None)

backwards_graph.structured_outputs = gradients_wrt_inputs

self._backward_graph_function = Function(

backwards_graph, attrs=backward_function_attr)

forward_function_attr = _parse_func_attrs({

BACKWARD_FUNCTION_ATTRIBUTE_NAME:

self._backward_graph_function._inference_function.name}) # pylint: disable=protected-access

forward_function_attr.update(self._attrs)

self._forward_function = _EagerDefinedFunction(

forward_function_name, self._func_graph, self._func_graph.inputs,

self._func_graph.outputs + backwards_graph_captures,

forward_function_attr)

def _backprop_call(self, args):

"""Calls the forward function and records the result on a tape.

(Only records results on a tape if the function has outputs)

Args:

args: All inputs to the function, including resolved captured inputs

Returns:

The call output.

"""

if self._backward_graph_function is None:

self._construct_backprop_function()

ctx = context.context()

if not self._gradient_name:

self._gradient_name = "PartitionedCall-%s" % ops.uid()

self._register_gradient(self._gradient_name)

with ops.get_default_graph().gradient_override_map(

{"PartitionedCall": self._gradient_name,

"StatefulPartitionedCall": self._gradient_name}):

outputs = self._forward_function.call(ctx, args)

if isinstance(outputs, ops.Operation) or outputs is None:

return outputs

# `real_outputs` are the actual outputs of the inference graph function;

# `side_outputs` are the intermediate Tensors that were added as outputs to

# the forward graph function so that we can compute its gradient.

real_outputs = outputs[:self._num_outputs]

skip_positions = [i for i, t in enumerate(real_outputs)

if not gradients_impl.IsTrainable(t)]

side_outputs = outputs[self._num_outputs:]

def backward_function(*args):

args = [a for i, a in enumerate(args)

if a is not None and i not in skip_positions]

return self._backward_graph_function(*(list(args) + side_outputs)) # pylint: disable=not-callable

tape.record_operation(self._forward_function.signature.name, real_outputs,

args, backward_function)

return self._build_call_outputs(real_outputs)

def _build_call_outputs(self, result):

"""Maps the fdef output list to actual output structure.

Args:

result: Output lists defined by FunctionDef.

Returns:

The actual call output.

"""

if self._func_graph.structured_outputs is None:

return result

# Use `nest.flatten` instead of `func_graph_module.flatten` in order to

# preserve any IndexedSlices in `self._func_graph.structured_outputs`.

outputs_list = nest.flatten(self._func_graph.structured_outputs)

j = 0

for i, o in enumerate(outputs_list):

if o is not None:

if isinstance(o, ops.IndexedSlices):

# Repack Tensors for IndexedSlices.

if o.dense_shape is not None:

outputs_list[i] = ops.IndexedSlices(

values=result[j],

indices=result[j + 1],

dense_shape=result[j + 2])

j += 3

else:

outputs_list[i] = ops.IndexedSlices(

values=result[j], indices=result[j + 1])

j += 2

else:

outputs_list[i] = result[j]

j += 1

ret = nest.pack_sequence_as(self._func_graph.structured_outputs,

outputs_list)

return ret

pywrap_tensorflow.RegisterType("Tensor", ops.Tensor)

pywrap_tensorflow.RegisterType("IndexedSlices", ops.IndexedSlices)

def _deterministic_dict_values(dictionary):

return tuple(dictionary[key] for key in sorted(dictionary))

class PolymorphicFunction(object):

"""Wrapper class for the graph functions defined for a Python function.

See the documentation for `defun` for more information on the semantics of

defined functions.

PolymorphicFunction class is thread-compatible meaning that minimal

usage of defuns (defining and calling) is thread-safe, but if users call other

methods or invoke the base `python_function` themselves, external

synchronization is necessary.

"""

def __init__(self,

python_function,

name,

input_signature=None,

attributes=None,

experimental_autograph=False):

"""Initializes a polymorphic function.

Args:

python_function: the function to be wrapped.

name: the name given to it.

input_signature: a possibly nested sequence of `TensorSpec` objects

specifying the input signature of this function. If `None`, a separate

function is instantiated for each inferred input signature.

attributes: dict, extra keyword arguments that will be added as attribute

of the function.

experimental_autograph: whether to use autograph to compile

`python_function`. See https://www.tensorflow.org/guide/autograph for

more information.

Raises:

ValueError: if `input_signature` is not None and the `python_function`'s

argspec has keyword arguments.

"""

if isinstance(python_function, functools.partial):

self._python_function = python_function.func

self._args_to_prepend = python_function.args or tuple()

self._kwargs_to_include = python_function.keywords or {}

else:

self._python_function = python_function

self._args_to_prepend = tuple()

self._kwargs_to_include = {}

self._name = name

self._experimental_autograph = experimental_autograph

self._function_cache = collections.OrderedDict()

self._function_attributes = attributes or {}

self._lock = threading.Lock()

# _descriptor_cache is a of instance of a class to an instance-specific

# PolymorphicFunction, used to make sure defun-decorated methods create

# different functions for each instance.

self._descriptor_cache = weakref.WeakKeyDictionary()

fullargspec = tf_inspect.getfullargspec(self._python_function)

if tf_inspect.ismethod(self._python_function):

# Remove `self`: default arguments shouldn't be matched to it.

args = fullargspec.args[1:]

else:

args = fullargspec.args

# A cache mapping from argument name to index, for canonicalizing

# arguments that are called in a keyword-like fashion.

self._args_to_indices = {arg: i for i, arg in enumerate(args)}

self._arg_names = args

self._vararg_name = fullargspec.varargs

# A cache mapping from arg index to default value, for canonicalization.

offset = len(args) - len(fullargspec.defaults or [])

self._arg_indices_to_default_values = {

offset + index: default

for index, default in enumerate(fullargspec.defaults or [])

}

self._default_values = fullargspec.defaults

self._default_values_start_index = offset

if input_signature is None:

self._input_signature = None

else:

if fullargspec.varkw is not None or fullargspec.kwonlyargs:

raise ValueError("Cannot define a TensorFlow function from a Python "

"function with keyword arguments when "

"input_signature is provided.")

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

raise TypeError("input_signature must be either a tuple or a "

"list, received " + str(type(input_signature)))

self._input_signature = tuple(input_signature)

self._flat_input_signature = tuple(nest.flatten(input_signature))

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

"""Calls a graph function specialized to the inputs."""

graph_function, inputs = self._maybe_define_function(args, kwargs)

return graph_function(*inputs)

@property

def python_function(self):

"""Returns the wrapped Python function."""

return self._python_function

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

"""Returns a `Function` object specialized to inputs and execution context.

`args` and `kwargs` are ignored if this `PolymorphicFunction` was created

with an `input_signature`.

Args:

*args: inputs to specialize on.

**kwargs: inputs to specialize on.

"""

if self._input_signature:

args, kwargs = None, None

graph_function, _ = self._maybe_define_function(args, kwargs)

return graph_function

def __get__(self, instance, owner):

"""Makes it possible to defun instance methods."""

del owner

# `instance` here is the instance that this `PolymorphicFunction` was

# accessed through; e.g., for

#

# class Foo(object):

#

# @function.defun

# def bar(self):

# ...

#

# foo = Foo()

# foo.bar() # `foo.bar` is a `PolymorphicFunction` instance

#

# then `instance` will be `foo` (and `owner` will be `Foo`). We create a

# new instance of PolymorphicFunction here to allow different instances each

# to create variables once, thereby allowing methods to be decorated with

# defun. Keeps a cache to avoid retracing the function every time the

# descriptor is accessed.

if instance not in self._descriptor_cache:

if instance is None:

return self

# If there is no instance-specific polymorphic func in the cache,

# we construct an instance-specific polymorphic function

# that uses a weak reference to the instance (so that the instance will

# be correctly gc'd).

# And finally add the wrapped function to the description cache

self._descriptor_cache[instance] = class_method_to_instance_method(

self, instance)

# Return the cached polymorphic function for the instance

return self._descriptor_cache[instance]

def _cache_key(self, args, kwargs):

"""Computes the cache key given inputs and execution context."""

if self._input_signature is None:

inputs = (args, kwargs) if kwargs else args

cache_key = pywrap_tensorflow.TFE_Py_EncodeArg(inputs)

else:

del args, kwargs

cache_key = self._flat_input_signature

ctx = context.context()

with ops.init_scope():

# The graph, or whether we're executing eagerly, should be a part of the

# cache key so we don't improperly capture tensors such as variables.

executing_eagerly = ctx.executing_eagerly()

execution_context = executing_eagerly or ops.get_default_graph()

# pylint: disable=protected-access

default_graph = ops.get_default_graph()

# TODO(b/117617952): The current distribution strategy will affect graph

# building (e.g. accessing different variables from different devices) and

# so requires retracing for each device.

uses_distribution_strategy = bool(

default_graph._distribution_strategy_stack)

if executing_eagerly:

colocation_stack = ()

uses_xla = ctx.device_spec.device_type == "TPU"

if uses_distribution_strategy or uses_xla:

device_functions = (pydev.merge_device(ctx.device_name),)

else:

device_functions = ()

else:

colocation_stack = tuple(default_graph._colocation_stack.peek_objs())

uses_xla = getattr(default_graph, "_xla_compile", False)

if (uses_distribution_strategy

or uses_xla

or func_graph_module.device_stack_has_callable(

default_graph._device_function_stack)):

# Putting the device in the cache key ensures that call-site device

# annotations are respected.

device_functions = tuple(default_graph._device_functions_outer_to_inner)

else:

device_functions = ()

# pylint: enable=protected-access

return (cache_key, execution_context, device_functions, colocation_stack,

uses_xla)

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

"""Canonicalizes `args` and `kwargs`.

Canonicalize the inputs to the Python function using its fullargspec. In

particular, we parse the varags and kwargs that this

`PolymorphicFunction` was called with into a tuple corresponding to the

Python function's positional (named) arguments and a dictionary

corresponding to its kwargs.

Args:

*args: The varargs this object was called with.

**kwargs: The keyword args this function was called with.

Returns:

A canonicalized ordering of the inputs.

Raises:

ValueError: If a keyword in `kwargs` cannot be matched with a positional

argument when an input signature is specified, or when the inputs

do not conform to the input signature.

"""

args = self._args_to_prepend + args

kwargs = dict(kwargs, **self._kwargs_to_include)

if not kwargs:

if self._default_values:

inputs = args + self._default_values[len(args) -

self._default_values_start_index:]

else:

inputs = args

else:

# Maps from index of arg to its corresponding value, according to `args`

# and `kwargs`; seeded with the default values for the named args that

# aren't in `args`.

arg_indices_to_values = {

index: default for index, default in six.iteritems(

self._arg_indices_to_default_values) if index >= len(args)

}

consumed_args = []

for arg, value in six.iteritems(kwargs):

index = self._args_to_indices.get(arg, None)

if index is not None:

arg_indices_to_values[index] = value

consumed_args.append(arg)

elif self._input_signature is not None:

raise ValueError("Cannot define a TensorFlow function from a Python "

"function with keyword arguments when "

"input_signature is provided.")

for arg in consumed_args:

# After this loop, `kwargs` will only contain true keyword arguments, as

# opposed to named arguments called in a keyword-like fashion.

kwargs.pop(arg)

inputs = args + _deterministic_dict_values(arg_indices_to_values)

flat_inputs = nest.flatten(inputs)

# Check for NumPy arrays in arguments and convert them to Tensors.

# TODO(nareshmodi): Skip ndarray conversion to tensor altogether, perhaps

# finding a way to store them directly in the cache key (currently not

# possible since ndarrays are not hashable).

need_packing = False

for index, value in enumerate(flat_inputs):

if type(value) == np.ndarray:

flat_inputs[index] = constant_op.constant(value)

need_packing = True

if need_packing:

inputs = nest.pack_sequence_as(structure=inputs,

flat_sequence=flat_inputs)

if self._input_signature is None:

return inputs, kwargs

else:

assert not kwargs

try:

nest.assert_same_structure(self._input_signature, inputs)

except (ValueError, TypeError):

raise ValueError("Structure of Python function inputs does not match "

"input_signature.")

if any(not pywrap_tensorflow.IsTensor(arg) for arg in flat_inputs):

raise ValueError("When input_signature is provided, all inputs to "

"the Python function must be Tensors.")

if any(not spec.is_compatible_with(other)

for spec, other in zip(self._flat_input_signature, flat_inputs)):

raise ValueError("Python inputs incompatible with input_signature: "

"inputs (%s), input_signature (%s)" %

(str(inputs), str(self._input_signature)))

return inputs, {}

def _maybe_define_function(self, args, kwargs):

"""Gets a function for these inputs, defining it if necessary.

`args` and `kwargs` can be None if this `PolymorphicFunction` was created

with an `input_signature`.

Args:

args: The varargs for the Python function.

kwargs: The keyword args for the Python function.

Returns:

A graph function corresponding to the input signature implied by args and

kwargs, as well as the inputs that the object should be called with.

Raises:

ValueError: If inputs are incompatible with the input signature.

TypeError: If the function inputs include non-hashable objects

"""

if self._input_signature is None or args is not None or kwargs is not None:

args, kwargs = self._canonicalize_function_inputs(*args, **kwargs)

cache_key = self._cache_key(args, kwargs)

with self._lock:

try:

graph_function = self._function_cache.get(cache_key, None)

except TypeError:

raise TypeError("Arguments supplied to `defun`-generated functions "

"must be hashable.")

if graph_function is None:

if self._input_signature is None:

arglen = len(args)

else:

arglen = len(self._input_signature)

arg_names = (

self._arg_names[:arglen]

+ [self._vararg_name] * (arglen - len(self._arg_names)))

graph_function = Function(

func_graph_module.func_graph_from_py_func(

self._name,

self._python_function,

args,

kwargs,

self._input_signature,

experimental_autograph=self._experimental_autograph,

arg_names=arg_names),

self._function_attributes)

self._function_cache[cache_key] = graph_function

return graph_function, [

t for t in nest.flatten((args, kwargs))

if isinstance(t, (ops.Tensor, resource_variable_ops.ResourceVariable))

]

def register_concrete(func):

"""Register a concrete function into the graph.

Args:

func: A graph function.

"""

graph = ops.get_default_graph()

# There are two situations for the actual call of a defun:

# 1. If none of the input args are resource variables or watch by any tape,

# it will run the _inference_function of concrete_func for forward pass, and

# the gradient will be generated by standard mechanism.

# 2. Otherwise, defun will create two functions, one for forward pass, and the

# backward pass will be created via tape.

# When registering the function, we put both cases into graph.

# pylint: disable=protected-access

func._inference_function.add_to_graph(graph)

if func._backward_graph_function is None:

func._construct_backprop_function()

forward_function = func._forward_function

backward_function = func._backward_graph_function._inference_function

forward_function.add_to_graph(graph)

backward_function.add_to_graph(graph)

# pylint: enable=protected-access

def register(func, *args, **kwargs):

"""Register a specialization of a PolymorphicFunction into the graph.

This won't actually call the function with the inputs, and only put the

function definition into graph. Register function with different input param

will result into multiple version of functions registered in graph.

Also, `args` and `kwargs` are ignored if this `PolymorphicFunction` was

created with an `input_signature`.

Args:

func: the PolymorphicFunction instance that generated by a @defun

*args: input arguments for the Python function.

**kwargs: input keyword arguments for the Python function.

Returns:

a `Function` object specialized to inputs and execution context.

Raises:

ValueError: When the input function is not a defun wrapped python function.

"""

if not isinstance(func, PolymorphicFunction):

raise ValueError("Only defun function is allowed to be registered. "

"Got type: %s" % type(func))

concrete_func = func.get_concrete_function(*args, **kwargs)

register_concrete(concrete_func)

return concrete_func

def validate_signature(signature):

if any(not isinstance(arg, tensor_spec.TensorSpec)

for arg in nest.flatten(signature)):

raise TypeError("Invalid input_signature %s; input_signature must be "

"a possibly nested sequence of TensorSpec objects.")