# 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 threading

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 dtypes as dtypes_module

from tensorflow.python.framework import ops

from tensorflow.python.framework import tensor_spec

from tensorflow.python.ops import array_ops

from tensorflow.python.ops import control_flow_ops

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.training import distribute

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

def create_substitute_placeholder(value, name, dtype=None):

"""Creates a placeholder for `value` and propagates shape info to it."""

# Note: setting ops.control_dependencies(None) ensures we always put

# capturing placeholders outside of any control flow context.

with ops.control_dependencies(None):

placeholder = graph_placeholder(

dtype=dtype or value.dtype, shape=value.shape, name=name)

if placeholder.dtype == dtypes_module.resource:

if isinstance(value, ops.EagerTensor):

handle_data = value._handle_data # pylint: disable=protected-access

else:

handle_data = resource_variable_ops.get_resource_handle_data(value)

if handle_data is not None and handle_data.is_set:

# pylint: disable=protected-access

pywrap_tensorflow.SetResourceHandleShapeAndType(

placeholder.graph._c_graph, placeholder._as_tf_output(),

handle_data.SerializeToString())

# pylint: enable=protected-access

# Ensure that shapes and dtypes are propagated.

shapes, types = zip(*[(pair.shape, pair.dtype)

for pair in handle_data.shape_and_type])

ranks = [len(s.dim) if not s.unknown_rank else -1 for s in shapes]

shapes = [[d.size for d in s.dim]

if not s.unknown_rank else None for s in shapes]

pywrap_tensorflow.TF_GraphSetOutputHandleShapesAndTypes_wrapper(

placeholder._op._graph._c_graph, # pylint: disable=protected-access

placeholder._as_tf_output(), # pylint: disable=protected-access

shapes, ranks, types)

return placeholder

def capture_value(tensor_map, value, dtype, name):

"""Capture a value from outside the function, to pass in as an extra arg."""

captured_tuple = tensor_map.get(ops.tensor_id(value), None)

if captured_tuple is None:

captured_value = create_substitute_placeholder(value, name=name,

dtype=dtype)

tensor_map[ops.tensor_id(value)] = (value, captured_value)

else:

captured_value = captured_tuple[1]

tape.record_operation("captured_value", [captured_value], [value],

lambda x: [x])

return captured_value

class CapturingGraph(ops.Graph):

"""Graph used when constructing eager functions."""

def __init__(self):

super(CapturingGraph, self).__init__()

self._building_function = True

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

self.captures = {}

# Map from resource tensor name to last op (in program order) which uses

# this tensor. Used to enforce that execution order matches program order

# for resource tensors.

self._last_op_using_resource_tensor = {}

# TODO(apassos) remove once the C API is used by default.

def _use_c_api_hack(self):

return True

def clear_resource_control_flow_state(self):

self._last_op_using_resource_tensor = {}

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

if isinstance(tensor, ops.EagerTensor):

if name is None:

name = str(ops.uid())

return capture_value(self.captures, tensor, tensor.dtype, name)

if tensor.graph is not self:

if name is None:

name = tensor.op.name

return capture_value(self.captures, tensor, tensor.dtype, name)

return tensor

def create_op(

self,

op_type,

inputs,

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

input_types=None,

name=None,

attrs=None,

op_def=None,

compute_shapes=True,

compute_device=True):

# 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

for i, inp in enumerate(inputs):

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

inp = ctxt.AddValue(inp)

inp = self.capture(inp)

inputs[i] = inp

return super(CapturingGraph, self).create_op(

op_type, inputs, dtypes, input_types, name, attrs, op_def,

compute_device=compute_device)

# pylint: disable=invalid-name

class HelperContext(object):

"""ControlFlowContext with a customizable AddOp method."""

def __init__(self, add_op_internal):

self._add_op_internal = add_op_internal

self._values = set() # control flow code sometimes updates this.

def _AddOpInternal(self, op):

self._add_op_internal(op)

@property

def outer_context(self):

return self._outer_context

def GetWhileContext(self):

if self._outer_context:

return self._outer_context.GetWhileContext()

def IsWhileContext(self):

return False

def IsCondContext(self):

return False

def IsXLAContext(self):

return False

def AddOp(self, op): # pylint: disable=invalid-name

self._AddOpInternal(op)

if self._outer_context:

self._outer_context.AddOp(op)

def AddName(self, _):

pass

def AddInnerOp(self, op):

self._AddOpInternal(op)

if self._outer_context:

self._outer_context.AddInnerOp(op)

def AddValue(self, val):

if self._outer_context:

return self._outer_context.AddValue(val)

else:

return val

def EnterGradientColocation(self, op, gradient_uid):

"""Start building a gradient colocated with an op."""

if self._outer_context:

self._outer_context.EnterGradientColocation(op, gradient_uid)

def ExitGradientColocation(self, op, gradient_uid):

"""Start building a gradient colocated with an op."""

if self._outer_context:

self._outer_context.ExitGradientColocation(op, gradient_uid)

def __enter__(self):

# pylint: disable=protected-access

self._g = ops.get_default_graph()

self._outer_context = self._g._get_control_flow_context()

self._g._set_control_flow_context(self)

self._nested_contexts = (

self._outer_context._nested_contexts

if self._outer_context is not None else None)

# pylint: enable=protected-access

def __exit__(self, *_):

self._g._set_control_flow_context(self._outer_context) # pylint: disable=protected-access

# pylint: enable=invalid-name

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

def _register(fn):

"""Registers the function `fn`."""

context.context().add_function(fn)

_xla_compile_attr = "_XlaCompile"

# 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, operations, 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

operations: list of Operation; the subset of operations in the graph

which will be 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

"""

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)

self._xla_compile = _xla_compile_attr in attrs

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

if context.executing_eagerly():

_register(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.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, output_shapes):

"""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.

output_shapes: shapes to which outputs should be set; ignored when

executing eagerly.

Returns:

The outputs of the function call.

"""

executing_eagerly = ctx.executing_eagerly()

xla_compile = self._xla_compile or (executing_eagerly and

ctx.device_spec.device_type == "TPU")

if xla_compile:

# 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`.

outputs = functional_ops.partitioned_call(

args=args,

f=self,

tout=self._output_types,

executing_eagerly=executing_eagerly)

if executing_eagerly:

return outputs

else:

for i, shape in enumerate(output_shapes):

outputs[i].set_shape(shape)

return outputs

def _map_sequence_obj_to_idx(sequence):

"""Maps objs in the sequence from id(obj) to sequence index."""

return {id(x): i for i, x in enumerate(sequence)}

def _flatten(sequence):

"""A wrapper around `nest.flatten` that also unpacks `IndexedSlices`."""

# TODO(akshayka): Support `SparseTensor` in a similar fashion.

flat_sequence = nest.flatten(sequence)

outputs = []

for item in flat_sequence:

if isinstance(item, ops.IndexedSlices):

if item.dense_shape is not None:

outputs.extend([item.values, item.indices, item.dense_shape])

else:

outputs.extend([item.values, item.indices])

else:

outputs.append(item)

return outputs

# TODO(akshayka): Perhaps rename to something more appropriate.

class GraphModeFunction(object):

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

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

is differentiable under `tf.GradientTape` objects.

"""

def __init__(self,

name,

input_placeholders,

extra_inputs,

graph,

operations,

outputs,

python_func_outputs,

output_shapes,

variables=None,

attrs=None):

"""Initialize a GraphModeFunction.

Args:

name: str the name of the created function

input_placeholders: list of placeholder values (tensors) to feed when

calling the wrapped function.

extra_inputs: Tensor inputs this function definition closed over which

are passed as arguments. Need to track so gradients are supported

correctly.

graph: the Graph from which the operations will be pulled. Used as

a context when computing gradients.

operations: the subset of Operations in the graph used in the function

definition.

outputs: a flat list of the Tensors in the graph used as outputs to the

function

python_func_outputs: a possibly nested python object which will be

returned by this function. The Tensors in this structure will be

replaced by their corresponding values in outputs. Note that this

structure might contain Python `None`s.

output_shapes: List of shapes of all tensors in outputs

variables: (optional) List of variables to watch during function

execution.

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

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

definition.

"""

self._attrs = attrs or {}

defined_function = _EagerDefinedFunction(

name, graph, operations, input_placeholders, outputs, self._attrs)

if len(input_placeholders) != len(defined_function.signature.input_arg):

raise ValueError("Internal error: invalid lengths. %s %s" % (

len(input_placeholders), len(defined_function.signature.input_arg)))

self._input_placeholders = input_placeholders

self._extra_inputs = list(extra_inputs)

self._graph = graph

self._backward_function = None

self._func_name = name

self._function_def = defined_function

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

self._python_func_outputs = python_func_outputs

self._python_returns = [python_func_outputs] if isinstance(

python_func_outputs,

(ops.Tensor, type(None))) else _flatten(python_func_outputs)

self._output_shapes = output_shapes

self._variables = variables if variables is not None else []

# Find the variables that are components of something distributed and

# put them into a {handle_tensor -> distributed variable object} map.

self._distributed_variables = {}

strategy = distribute.get_distribution_strategy()

for variable in self._variables:

# If variable is not distributed, unwrap returns [variable].

component_variables = strategy.unwrap(variable)

# Only add to the dictionary when the variable is actually distributed,

# i.e. more than one component or the component is different from the

# variable itself. component_variables cannot be empty.

if (len(component_variables) > 1 or component_variables[0] != variable):

for component_variable in component_variables:

self._distributed_variables[component_variable.handle] = variable

@property

def variables(self):

return self._variables

def _construct_backprop_function(self):

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

filtered_outputs = [x for x in self._python_returns if x is not None]

backwards_graph = CapturingGraph()

backwards_graph._graph_key = self._graph._graph_key # pylint: disable=protected-access

for collection in self._graph.collections:

backwards_graph.get_collection_ref(

collection)[:] = self._graph.get_collection(collection)

backwards_graph.seed = self._graph.seed

with backwards_graph.as_default():

self._out_grad_placeholders = [

graph_placeholder(x.dtype, x.shape) for x in filtered_outputs]

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

filtered_outputs,

self._input_placeholders,

grad_ys=self._out_grad_placeholders,

src_graph=self._graph)

backward_outputs = tuple(

grad for grad in _flatten(in_gradients) if grad is not None)

output_shapes = tuple(grad.shape for grad in backward_outputs)

captures = backwards_graph.captures

ids = list(sorted(captures.keys()))

if ids:

extra_inputs, extra_placeholders = zip(*[captures[x] for x in ids])

else:

extra_inputs = []

extra_placeholders = []

forward_name = _forward_name(self._func_name)

# Note: we cannot have placeholder ops in the graph or the TPU compilation

# pass fails.

placeholder_ops = set([y.op for y in self._input_placeholders])

function_ops = [x for x in self._graph.get_operations()

if x not in placeholder_ops]

self._forward_fdef = _EagerDefinedFunction(

forward_name, self._graph, function_ops,

self._input_placeholders, filtered_outputs + list(extra_inputs),

self._attrs)

all_inputs = self._out_grad_placeholders + list(extra_placeholders)

# Excluding input ops from the body as we do not intend to execute these

# operations when the function is executed.

all_ignored_ops = frozenset(x.op for x in all_inputs)

# Enforce a deterministic order of operations in the generated graph. This

# means rerunning the function-defining code will always define the same

# function, which is useful if we serialize this etc.

function_def_ops = tuple(x

for x in sorted(backwards_graph.get_operations(),

key=lambda x: x.name)

if x not in all_ignored_ops)

bname = _backward_name(self._func_name)

self._backward_function = GraphModeFunction(

bname, all_inputs, [], backwards_graph, function_def_ops,

backward_outputs, in_gradients, output_shapes, attrs=self._attrs)

def _backprop_call(self, args):

"""Calls the wrapped 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 extra inputs

Returns:

The call output.

"""

ctx = context.context()

outputs = self._forward_fdef.call(ctx, args, self._output_shapes)

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]

side_outputs = outputs[self._num_outputs:]

def backward_function(*args):

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

tape.record_operation(

self._forward_fdef.signature.name,

real_outputs,

args,

backward_function)

return self._build_call_outputs(real_outputs)

@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?

outputs_list = nest.flatten(self._python_func_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._python_func_outputs, outputs_list)

@property

def output_dtypes(self):

return nest.map_structure(

lambda x: x.dtype if x is not None else None, self._python_func_outputs)

@property

def captured_inputs(self):

return self._extra_inputs

@property

def name(self):

"""Returns the name of the function in Eager-compatible format."""

return self._function_def.name.encode("utf-8")

def _resolve_extra_inputs(self):

"""Resolve captured distributed variables to their current values.

Some inputs can be distributed variables. Such variables yield a different

component (i.e. actual tf.Variable) variables depending on the context of

execution.

Returns:

a list of resolved extra input tensors.

"""

if self._distributed_variables:

# Loop over each extra_inputs and check if it corresponds to something

# distributed. If so, get its _distributed_container and fetch the

# component appropriate for the current execution context.

resolved_extra_inputs = self._extra_inputs[:]

for i, extra_input in enumerate(self._extra_inputs):

distributed_var = self._distributed_variables.get(extra_input, None)

if distributed_var is not None:

# distributed variables override __getattr__ and substitute the

# right component variable. In here, `distributed_var.handle`

# actually does the equivalent of

# distributed_var.get_current_component_var().handle.

resolved_extra_inputs[i] = distributed_var.handle

return resolved_extra_inputs

return self._extra_inputs

def __call__(self, *args):

"""Executes the passed function in eager mode."""

for v in self._variables:

if v.trainable:

tape.watch_variable(v)

resolved_extra_inputs = self._resolve_extra_inputs()

tensor_inputs = [x for x in nest.flatten(args) if isinstance(x, ops.Tensor)]

args = tensor_inputs + resolved_extra_inputs

if tape.should_record(tensor_inputs) or tape.should_record(

resolved_extra_inputs):

if self._backward_function is None:

self._construct_backprop_function()

return self._backprop_call(args)

ctx = context.context()

outputs = self._function_def.call(ctx, args, self._output_shapes)

return self._build_call_outputs(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._python_func_outputs is None:

return result

# Use `nest.flatten` instead of `_flatten` in order to preserve any

# IndexedSlices in `self._python_func_outputs`.

outputs_list = nest.flatten(self._python_func_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._python_func_outputs, outputs_list)

return ret

def _get_defun_inputs_from_signature(signature):

"""Maps a signature to graph-construction inputs."""

function_inputs = [

graph_placeholder(spec.dtype, spec.shape)

for spec in nest.flatten(signature)

]

return nest.pack_sequence_as(signature, function_inputs)

def _get_defun_inputs_from_args(args):

"""Maps python function args to graph-construction inputs."""

function_inputs = [

graph_placeholder(arg.dtype, arg.shape) if isinstance(arg, ops.Tensor)

else arg for arg in nest.flatten(args)

]

return nest.pack_sequence_as(args, function_inputs)

def _trace_and_define_function(name, python_func, compiled, args, kwds,

signature=None):

"""Defines and returns graph-mode version of `python_func`.

Args:

name: an identifier for the function.

python_func: the Python function to trace.

compiled: whether the graph function should be compiled through XLA.

args: the positional args with which the Python function should be called;

ignored if a signature is provided.

kwds: the keyword args with which the Python function should be called;

ignored if a signature is provided.

signature: a possibly nested sequence of `TensorSpecs` specifying the shapes

and dtypes of the arguments. When a signature is provided, `args` and

`kwds` are ignored, and `python_func` is traced with Tensors conforming

to `signature`. If `None`, the shapes and dtypes are inferred from the

inputs.

Returns:

A GraphModeFunction.

"""

graph_key = ops.get_default_graph()._graph_key # pylint: disable=protected-access

func_graph = CapturingGraph()

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

# optimizers.

func_graph._graph_key = graph_key # pylint: disable=protected-access

# Copy the graph collections to ensure summaries and other things work. This

# lets the function access (but not mutate) collections of the containing

# graph, such as the global step and the summary writer collections.

curr_graph = ops.get_default_graph()

for collection in curr_graph.collections:

func_graph.get_collection_ref(collection)[:] = curr_graph.get_collection(

collection)

if context.executing_eagerly():

func_graph.seed = context.global_seed()

else:

func_graph.seed = curr_graph.seed

with func_graph.as_default(), AutomaticControlDependencies() as a:

if signature is None:

func_args = _get_defun_inputs_from_args(args)

func_kwds = _get_defun_inputs_from_args(kwds)

else:

func_args = _get_defun_inputs_from_signature(signature)

func_kwds = {}

# Variables to help check whether mutation happens in calling the function

# Copy the recursive list, tuple and map structure, but not base objects

func_args_before = nest.pack_sequence_as(func_args, nest.flatten(func_args))

func_kwds_before = nest.pack_sequence_as(func_kwds, nest.flatten(func_kwds))

def convert(x):

if x is None:

return None

x = ops.convert_to_tensor_or_indexed_slices(x)

x = a.mark_as_return(x)

return x

this_tape = tape.push_new_tape()

try:

func_outputs = python_func(*func_args, **func_kwds)

func_outputs = nest.map_structure(convert, func_outputs)

def check_mutation(n1, n2):

"""Check if two list of arguments are exactly the same."""

errmsg = ("Function to be traced should not modify structure of input "

"arguments. Check if your function has list and dictionary "

"operations that alter input arguments, "

"such as `list.pop`, `list.append`")

try:

nest.assert_same_structure(n1, n2)

except ValueError:

raise ValueError(errmsg)

for arg1, arg2 in zip(nest.flatten(n1), nest.flatten(n2)):

if arg1 is not arg2:

raise ValueError(errmsg)

check_mutation(func_args_before, func_args)

check_mutation(func_kwds_before, func_kwds)

finally:

tape.pop_tape(this_tape)

variables = list(this_tape.watched_variables())

# Some variables captured by the tape can come from a DistributedValue.

# At call time, DistributedValue can return another variable (e.g. if

# the function is run on a different device). Thus, instead of storing

# the specific captured variable, we replace it with its distributed

# container.

strategy = distribute.get_distribution_strategy()

for i, variable in enumerate(variables):

# If variable is not distributed value_container returns itself.

variables[i] = strategy.value_container(variable)

# Returning a closed-over tensor as an output does not trigger a

# call to convert_to_tensor, so we manually capture all such tensors.

outputs_list = _flatten(func_outputs)

func_def_outputs = [

func_graph.capture(x) for x in outputs_list

if x is not None

]

captures = func_graph.captures

ids = list(sorted(captures.keys()))

if ids:

extra_inputs, extra_placeholders = zip(* [captures[x] for x in ids])

else:

extra_inputs = []

extra_placeholders = []

output_shapes = tuple(

x.shape if isinstance(x, ops.Tensor) else None

for x in func_def_outputs)

# Note: `nest.flatten` sorts by keys, as does `_deterministic_dict_values`.

flat_inputs = [

x for x in nest.flatten(func_args) + nest.flatten(func_kwds)

if isinstance(x, ops.Tensor)

]

all_inputs = flat_inputs + list(extra_placeholders)

all_ignored_ops = frozenset(x.op for x in all_inputs)

fname = _inference_name(name)

operations = tuple(x for x in func_graph.get_operations()

if x not in all_ignored_ops)

# Register any other functions defined in the graph

# TODO(ashankar): Oh lord, forgive me for this lint travesty.

if context.executing_eagerly():

for f in func_graph._functions.values(): # pylint: disable=protected-access

# TODO(ashankar): What about the gradient registry?

_register(f._c_func.func) # pylint: disable=protected-access

attrs = {}

if compiled:

attrs[_xla_compile_attr] = attr_value_pb2.AttrValue(b=True)

return GraphModeFunction(

fname, all_inputs, extra_inputs, func_graph, operations, func_def_outputs,

func_outputs, output_shapes, variables, attrs)

_TensorType = collections.namedtuple("_TensorType", ["dtype", "shape"])

def _encode_arg(arg):

"""A canonical representation for this argument, for use in a cache key."""

# `defun` uses dtypes and shapes instead of `Tensors` as cache keys. Dtypes

# are used because TensorFlow graphs are not parametric w.r.t. dtypes. Shapes

# are used for both performance reasons, as much TensorFlow code specializes

# on known shapes to produce slimmer graphs, and correctness, as some

# high-level APIs require shapes to be fully-known.

#

# TODO(akshayka): Add support for sparse tensors.

#

# pylint: disable=protected-access

if isinstance(arg, ops.Tensor):

return _TensorType(arg.dtype, arg._shape_tuple())

elif isinstance(arg, ops.IndexedSlices):

if arg.dense_shape is not None:

return tuple([

_TensorType(arg.values.dtype, arg.values._shape_tuple()),

_TensorType(arg.indices.dtype, arg.indices._shape_tuple()),

_TensorType(arg.dense_shape.dtype, arg.dense_shape._shape_tuple()),

])

else:

return tuple([

_TensorType(arg.values.dtype, arg.values._shape_tuple()),

_TensorType(arg.indices.dtype, arg.indices._shape_tuple()),

])

elif isinstance(arg, np.ndarray):

tensor = ops.convert_to_tensor(arg)

return _TensorType(tensor.dtype, tensor._shape_tuple())

# pylint: enable=protected-access

elif isinstance(arg, (list, tuple)):

return tuple([_encode_arg(elem) for elem in arg])

elif isinstance(arg, dict):

return tuple(

(_encode_arg(key), _encode_arg(arg[key])) for key in sorted(arg))

else:

return arg

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,

compiled=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.

compiled: if True, the framework will attempt to compile func with XLA.

Raises:

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

argspec has keyword arguments.

TypeError: if `input_signature` contains anything other than

`TensorSpec` objects, or (if not None) is anything other than a tuple or

list.

"""

if isinstance(python_function, functools.partial):

self._python_function = python_function.func

self._args_to_prepend = python_function.args or tuple()

self._kwds_to_include = python_function.keywords or {}

else:

self._python_function = python_function

self._args_to_prepend = tuple()

self._kwds_to_include = {}

self._name = name

self._compiled = compiled

self._arguments_to_functions = {}

self._variables = []

self._lock = threading.Lock()

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

# 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 [])

}

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

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

for arg in self._flat_input_signature):

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

"a possibly nested sequence of TensorSpec objects.")

def __get__(self, instance, owner):

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

del owner

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