"""Rewrite input signature to be hashable.

# TODO(slebedev): consider using nest.

if isinstance(elem, tuple):

return tuple(map(_make_input_signature_hashable, elem))

try:

hash(elem)

except TypeError:

# TFE_Py_EncodeArg weakrefs arguments it does not recognize, and we expect

# all recognized types to be hashable.

assert isinstance(elem, weakref.ReferenceType)

v = elem()

if resource_variable_ops.is_resource_variable(v):

# We special case variables here to use unique_id as the cache key. This

# ensures we have to retrace whenever a different variable is passed in.

# This is needed to support cases where the user may use the id of a

# variable in the function perhaps as a lookup in a dictionary.

#

# This choice leads to more retracing when we could have possibly used the

# shape and dtype instead. However, we expect the number of variables in a

# program to be bounded, and correspondingly the number of retraces.

#

# Note we also include the class name to avoid collisions with strings.

return v.__class__, v._unique_id # pylint: disable=protected-access

if _is_ndarray(v):

# Numpy arrays are not hashable, but when calling functions we treat them

# in the same way as tf.Tensors.

if not hasattr(v, "shape") or not hasattr(v, "dtype"):

# TODO(tomhennigan) De-dup with _as_ndarray in _convert_numpy_inputs.

v = _as_ndarray(v)

return tensor_spec.TensorSpec(v.shape, v.dtype)

raise ValueError("Arguments to a tf.function must be Tensors, Variables, "

"or hashable Python objects (or nested structures of "

"these types).\nGot type: %s" % type(v).__name__)

"input_signature",

"parent_graph",

"device_functions",

"colocation_stack",

"in_cross_replica_context",

"xla_context_id",

"""Returns a TypeSpec for `x`, or `None` if `x` doesn't have a TensorSpec."""

if isinstance(x, ops.Tensor):

return tensor_spec.TensorSpec.from_tensor(x)

elif isinstance(x, type_spec.TypeSpec):

return x

elif isinstance(x, composite_tensor.CompositeTensor):

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

else:

"""Returns true if TypeSpec `b` is a subset of type `a` (or if a is None.)"""

if a is None:

return True

else:

"""Returns a shape-relaxed TypeSpec for x (if composite) or x (if not)."""

if isinstance(x, composite_tensor.CompositeTensor):

# pylint: disable=protected-access

return x._type_spec._with_tensor_ranks_only()

else:

"""Find a `TensorShape` that is compatible with both `x` and `y`."""

if x is None != y is None:

raise RuntimeError(

"Cannot find a common shape when LHS shape is None but RHS shape "

"is not (or vice versa): %s vs. %s" % (x, y))

if x is None:

return None # The associated input was not a Tensor, no shape generated.

if not isinstance(x, tensor_shape.TensorShape):

raise TypeError("Expected x to be a TensorShape but saw %s" % (x,))

if not isinstance(y, tensor_shape.TensorShape):

raise TypeError("Expected y to be a TensorShape but saw %s" % (y,))

if x.rank != y.rank or x.rank is None:

return tensor_shape.TensorShape(None)

dims = []

for dim_x, dim_y in zip(x.dims, y.dims):

if (dim_x != dim_y

or tensor_shape.dimension_value(dim_x) is None

or tensor_shape.dimension_value(dim_y) is None):

dims.append(None)

else:

dims.append(tensor_shape.dimension_value(dim_x))

structure2,

check_values=False):

"""Check two structures for equality, optionally of types and of values."""

try:

nest.assert_same_structure(structure1, structure2, expand_composites=True)

except (ValueError, TypeError):

return False

if check_values:

flattened1 = nest.flatten(structure1, expand_composites=True)

flattened2 = nest.flatten(structure2, expand_composites=True)

# First check the types to avoid AttributeErrors.

if any(type(f1) != type(f2) for f1, f2 in zip(flattened1, flattened2)):

return False

return flattened1 == flattened2

"""Convert the keyword arguments into function_def attributes.

attrs = {}

for key, value in attributes.items():

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, six.text_type)):

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

"""Context Manager that interpolates the exception from 'top_level_func'."""

def __init__(self, top_level_func):

self._func = top_level_func

def __enter__(self):

pass

def __exit__(self, typ, exc, tb):

if not exc or not isinstance(exc, errors.OpError):

return False

message = compat.as_text(exc.message)

_, tags = error_interpolation.parse_message(message)

g = None

func_stack = []

for t in tags:

if t.type == "function_node":

# TODO(mdan): Tests should cover this.

if t.name == compat.as_str(self._func.name):

g = self._func.graph

elif g:

next_func = g._get_function(t.name) # pylint: disable=protected-access

if next_func is not None and isinstance(next_func,

_EagerDefinedFunction):

g = next_func.graph

if g:

func_stack.append(g.name)

else:

func_stack.append("<unknown>")

if g:

message = error_interpolation.interpolate(message, g)

message += "\n\nFunction call stack:\n"

message += " -> ".join(func_stack)

message += "\n"

exc._message = message # pylint: disable=protected-access

"""Add a callback function for Function creation.

"""Remove an already-added function callback.

"""Returns the XLAControlFlowContext, which exists inside a tpu.rewrite()."""

graph = ops.get_default_graph()

while graph is not None:

# pylint: disable=protected-access

context_ = graph._get_control_flow_context()

# pylint: enable=protected-access

while context_ is not None:

if isinstance(context_, control_flow_ops.XLAControlFlowContext):

return context_

context_ = context_.outer_context

# This may be a FuncGraph due to defuns or v2 control flow. We need to

# find the original graph with the XLAControlFlowContext.

graph = getattr(graph, "outer_graph", None)

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

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

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

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

"""Initializes an eager defined function.

input_ops = set(arg.op for arg in inputs)

operations = [op for op in graph.get_operations() if op not in input_ops]

graph_output_names = graph._output_names # pylint: disable=protected-access

if (graph_output_names is not None and

all(ops.tensor_id(t) in graph_output_names for t in outputs)):

output_names = [

compat.as_bytes(graph_output_names[ops.tensor_id(t)]) for t in outputs

]

if len(set(output_names)) != len(output_names):

# There are duplicate names for some reason, probably an invalid

# signature. Revert to auto-naming.

output_names = []

else:

output_names = []

fn = pywrap_tf_session.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

output_names,

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

[], # control_output_names

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_tf_session.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_tf_session.TF_FunctionToFunctionDef(fn, buffer_)

proto_data = pywrap_tf_session.TF_GetBuffer(buffer_)

function_def = function_pb2.FunctionDef()

function_def.ParseFromString(compat.as_bytes(proto_data))

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

with ops.init_scope():

if context.executing_eagerly():

context.ensure_initialized()

context.add_function(fn)

self._function_deleter = _EagerDefinedFunctionDeleter(self.name)

self._registered_on_context = True

self.definition = function_def

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._control_captures = graph.control_captures

# Shallow copy outputs since ConcreteFunction may mutate it.

self._func_graph_outputs = list(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._is_stateful) # pylint: disable=protected-access

for function_callback in _function_callbacks:

function_callback(self)

def add_to_graph(self, g=None):

# pylint: disable=protected-access

if not g and context.executing_eagerly():

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

else:

if not g._is_function(self.name):

g._add_function(self)

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

if not g._is_function(f.name):

g._add_function(f)

# pylint: enable=protected-access

@property

def name(self):

return self._name

@property

def stateful_ops(self):

return self._stateful_ops

def call(self, ctx, args, cancellation_manager=None):

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

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

raise ValueError(

"Arguments and signature arguments do not match. "

"got: %s, expected: %s " %

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

function_call_options = ctx.function_call_options

if function_call_options.config_proto_serialized is None:

config = function_utils.get_disabled_rewriter_config()

else:

config = function_call_options.config_proto_serialized

executor_type = function_call_options.executor_type or ""

executing_eagerly = ctx.executing_eagerly()

attrs = ("executor_type", executor_type, "config_proto", config)

if executing_eagerly:

with _InterpolateFunctionError(self):

if cancellation_manager is None:

outputs = execute.execute(

str(self.signature.name),

num_outputs=self._num_outputs,

inputs=args,

attrs=attrs,

ctx=ctx)

else:

outputs = execute.execute_with_cancellation(

str(self.signature.name),

num_outputs=self._num_outputs,

inputs=args,

attrs=attrs,

ctx=ctx,

cancellation_manager=cancellation_manager)

# Replace empty list with None

outputs = outputs or None

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

with _InterpolateFunctionError(self):

with ops.control_dependencies(self._control_captures):

# The caller must use record_operation to record this operation in the

# eager case, so we enforce the same requirement for the non-eager

# case by explicitly pausing recording. We don't have a gradient

# registered for PartitionedCall, so recording this operation confuses

# forwardprop code (GradientTape manages to ignore it).

with tape.stop_recording():

outputs = functional_ops.partitioned_call(

args=args,

f=self,

tout=self._output_types,

executing_eagerly=executing_eagerly,

config=config,

executor_type=executor_type)

if executing_eagerly:

return outputs

else:

# TODO(b/128924522): This additional set_shape should not be

# necessary. ShapeRefiner likely needs to inspect handle_data. Remove this

# once that's done.

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

"""Caches forward/backward functions with a delayed forward rewrite."""

def __init__(self, func_graph, attrs, func_graph_deleter):

"""Construct an inference function and initialize caches."""

# A map from the number of forward function outputs with accepted gradients

# to forward and backward functions, used to cache non-tape backward

# function generation.

self._cached_function_pairs = {}

self._func_graph = func_graph

self._inference_function = _EagerDefinedFunction(

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

self._func_graph.inputs, self._func_graph.outputs, attrs)

self._attrs = attrs

self._gradient_name = None

# Note that the FuncGraph is mutated later, so we need to inspect it now to

# figure out the user-specified outputs of the inference function.

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

self._func_graph_deleter = func_graph_deleter

def forward_backward(self, num_doutputs=None):

"""A possibly-cached pair of forward and backward functions."""

if num_doutputs is None:

num_doutputs = self._num_inference_outputs

forward_backward = self._cached_function_pairs.get(num_doutputs)

if forward_backward is not None:

return forward_backward

forward, backward = self._construct_forward_backward(num_doutputs)

self._cached_function_pairs[num_doutputs] = (forward, backward)

return forward, backward

def _construct_forward_backward(self, num_doutputs):

"""Constructs a pair of forward and backward functions.

trainable_outputs = [

output for output in self._func_graph.outputs[:num_doutputs]

if backprop_util.IsTrainable(output)]

signature = []

for t in trainable_outputs:

signature.append(

tensor_spec.TensorSpec(*default_gradient.shape_and_dtype(t)))

def _backprop_function(*grad_ys):

with ops.device(None):

return gradients_util._GradientsHelper( # pylint: disable=protected-access

trainable_outputs,

self._func_graph.inputs,

grad_ys=grad_ys,

src_graph=self._func_graph)

with self._func_graph.as_default():

backwards_graph = func_graph_module.FuncGraph(

_backward_name(self._func_graph.name))

func_graph_module.func_graph_from_py_func(

name=backwards_graph.name,

python_func=_backprop_function,

args=[], kwargs={},

signature=signature,

func_graph=backwards_graph)

backwards_graph_captures = backwards_graph.external_captures

captures_from_forward = [

c for c in backwards_graph_captures if

not isinstance(c, ops.EagerTensor) and c.graph is self._func_graph]

forward_function_name = _forward_name(self._func_graph.name)

# NB: forward and backward function have their "_implements"

# attribute set to None if it was present. This is because we don't

# support replacing those functions. If we do want for those functions

# to have implements function we need to provide a mechanism that

# would allow to identify all functions that call this one

# and trace and update their signatures as well. At the moment

# we disable this, until the tooling for doing this becomes available.

# See:

# https://github.com/tensorflow/community/blob/master/rfcs/20190610-standardizing-composite_ops.md#appendix-future-support-for-optimizing-gradient-functions

common_attributes = dict(self._attrs)

common_attributes.pop(IMPLEMENTS_ATTRIBUTE_NAME, None)

existing_outputs = object_identity.ObjectIdentitySet(

self._func_graph.outputs)

for capture in captures_from_forward:

if capture not in existing_outputs:

existing_outputs.add(capture)

self._func_graph.outputs.append(capture)

backward_function_attr = _parse_func_attrs(

{FORWARD_FUNCTION_ATTRIBUTE_NAME: forward_function_name})

backward_function_attr.update(common_attributes)

backward_function = ConcreteFunction(

backwards_graph, attrs=backward_function_attr)

forward_function_attr = _parse_func_attrs({

BACKWARD_FUNCTION_ATTRIBUTE_NAME:

backward_function.name})

forward_function_attr.update(common_attributes)

forward_function = _EagerDefinedFunction(

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

self._func_graph.outputs, forward_function_attr)

return forward_function, backward_function

def _rewrite_forward_and_call_backward(self, op, *doutputs):

"""Add outputs to the forward call and feed them to the grad function."""

forward_function, backwards_function = self.forward_backward(len(doutputs))

if not backwards_function.outputs:

return backwards_function.structured_outputs

forward_function.add_to_graph(op.graph)

# pylint: disable=protected-access

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

op._set_func_attr("f", forward_function.name)

op._set_type_list_attr("Tout", forward_function._output_types)

op._add_outputs(

forward_function._output_types[len(op.outputs):],

forward_function._output_shapes[len(op.outputs):])

for i in range(len(op.outputs)):

func_graph_output = forward_function._func_graph_outputs[i]

custom_gradient.copy_handle_data(func_graph_output, op.outputs[i])

# pylint: enable=protected-access

capture_mapping = dict(

zip((ops.tensor_id(t) for t in self._func_graph.outputs), op.outputs))

remapped_captures = [

capture_mapping.get(ops.tensor_id(capture), capture)

for capture in backwards_function.captured_inputs

]

# Replace Nones with zeros since we're calling a graph function which

# expects numeric inputs.

cleaned_doutputs = []

for doutput, placeholder in zip(doutputs, self._func_graph.outputs):

if backprop_util.IsTrainable(placeholder):

if isinstance(doutput, ops.IndexedSlices):

# Gradient passed to a backward ConcreteFunction must be tf.Tensor,

# so we convert tf.IndexedSlices to tf.Tensor.

cleaned_doutputs.append(ops.convert_to_tensor(doutput))

elif doutput is not None:

cleaned_doutputs.append(doutput)

else:

cleaned_doutputs.append(default_gradient.zeros_like(placeholder))

# Compute the gradients using the side outputs

return backwards_function._call_flat( # pylint: disable=protected-access

cleaned_doutputs, remapped_captures)

def get_gradient_function(self):

"""Returns gradient function.

return self._rewrite_forward_and_call_backward

def forward(self, inference_args=None, input_tangents=None):

"""A forward function with only user-specified outputs.

del inference_args # unused

if input_tangents:

# This class does not support special-cased forwardprop. The arguments are

# here for compatibility with _TapeGradientFunctions.

raise AssertionError(

"Internal error: unexpectedly got forwardprop information in a class "

"that does not support forwardprop.")

return self._inference_function

def _backward(self, outputs):

"""Fetch a backward function for `outputs` from the forward function."""

def _backward_function(*args):

call_op = outputs[0].op

return self._rewrite_forward_and_call_backward(call_op, *args)

return _backward_function, outputs

def record(self, flat_outputs, inference_args, input_tangents):

"""Record the function call operation.

backward_function, to_record = self._backward(flat_outputs)

tape.record_operation(self._inference_function.signature.name,

to_record, inference_args + input_tangents,

"_ForwardWrapper", (

# The wrapper Graph.

"graph",

# A flat list of non-tangent Tensor outputs from the wrapped forward

# function.

"outputs",

# Indices for output tangents, same format as

# forwardprop_util.pack_tangents.

"output_indices",

# A flat list of tangents for `outputs`.

"""Caches forward and backward functions compatible with eager gradients.

def __init__(self, func_graph, attrs, func_graph_deleter,

forwardprop_input_indices, delayed_rewrite_functions,

need_gradients_for_jvps):

self._func_graph = func_graph

self._forward_graph = None

self._attrs = attrs

self._forward = None

self._backward = None

self._num_outputs = len(func_graph.outputs)

self._func_graph_deleter = func_graph_deleter

self._forwardprop_input_indices = forwardprop_input_indices

self._forwardprop_output_indices = None

self._num_forwardprop_outputs = 0

self._num_inference_outputs = len(func_graph.outputs)

self._num_trainable_inference_outputs = len(

[t for t in func_graph.outputs if backprop_util.IsTrainable(t)])

self._delayed_rewrite_functions = delayed_rewrite_functions

self._need_gradients_for_jvps = need_gradients_for_jvps

def _build_functions_for_outputs(

self, outputs, inference_args, input_tangents):

"""Forward+backward functions where the backward function sees `outputs`."""

# First figure out which of `outputs` are trainable. We'll accept gradients

# for each of these in the backward function.

handles_to_variables = self._func_graph.variable_captures

trainable_outputs = []

trainable_indices = []

for index, output in enumerate(outputs):

if backprop_util.IsTrainable(output):

# Swap in the Variable object for resource handles if we can so

# sparse gradients work.

output = handles_to_variables.get(id(output), output)

trainable_outputs.append(output)

trainable_indices.append(index)

backwards_graph = func_graph_module.FuncGraph(

_backward_name(self._func_graph.name))

with backwards_graph.as_default():

gradients_wrt_outputs = []

for output in trainable_outputs:

gradient_shape, gradient_dtype = default_gradient.shape_and_dtype(

output)

gradients_wrt_outputs.append(

graph_placeholder(gradient_dtype, gradient_shape))

with ops.device(None):

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

trainable_outputs,

self._func_graph.inputs,

grad_ys=gradients_wrt_outputs,

src_graph=self._func_graph)

captures_from_forward = [

c for c in backwards_graph.external_captures

if not isinstance(c, ops.EagerTensor) and c.graph is self._func_graph

]

existing_outputs = object_identity.ObjectIdentitySet(

self._func_graph.outputs)

for capture in captures_from_forward:

if capture not in existing_outputs:

existing_outputs.add(capture)

self._func_graph.outputs.append(capture)

forward_function_name = _forward_name(self._func_graph.name)

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

# `backward_function` correspond to outputs (including

# side outputs) of `self._tape_forward_function`.

backwards_graph.inputs = (

gradients_wrt_outputs + backwards_graph.internal_captures)

backwards_graph.outputs.extend(

grad

for grad in nest.flatten(gradients_wrt_inputs, expand_composites=True)

if grad is not None)

backwards_graph.structured_outputs = gradients_wrt_inputs

backward_function = ConcreteFunction(

backwards_graph, attrs=backward_function_attr)

forward_function_attr = _parse_func_attrs({

BACKWARD_FUNCTION_ATTRIBUTE_NAME:

backward_function.name})

forward_function_attr.update(self._attrs)

forward_function = _EagerDefinedFunction(

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

self._func_graph.outputs,

forward_function_attr)

if not input_tangents:

# There is no need to special-case forwardprop, so we can return the

# forward+backward pair we've created without further wrapping.

return (forward_function, self._func_graph, backward_function,

# No forwardprop outputs.

None, 0)

forward_wrapper = self._wrap_forward_function_with_jvps(

forward_function, backward_function, inference_args, input_tangents)

(wrapped_backwards_graph,

forward_wrapper) = self._wrap_backward_function_with_jvp_backprop(

backward_function, gradients_wrt_outputs, forward_wrapper)

# Now that we've added new captures, we need to make sure forward outputs

# are in the same order the backward function expects them to be in:

# [inference outputs] + [jvps] + [side outputs] + [captures].

forward_wrapper = self._shuffle_forward_outputs(forward_wrapper)

wrapped_forward_function = _EagerDefinedFunction(

_forward_name(self._func_graph.name), forward_wrapper.graph,

forward_wrapper.graph.inputs, forward_wrapper.graph.outputs,

forward_function_attr)

wrapped_backward_function = ConcreteFunction(

wrapped_backwards_graph, attrs=backward_function_attr)

if (len(inference_args) + len(input_tangents)

!= len(forward_wrapper.graph.inputs)):

raise AssertionError(

("Internal error: the forward graph had {} inputs, but we expected"

" {} ({} inference inputs and {} input tangents)")

.format(len(len(forward_wrapper.graph.inputs)),

len(inference_args) + len(input_tangents),

len(inference_args), len(input_tangents)))

return (wrapped_forward_function, forward_wrapper.graph,

wrapped_backward_function, forward_wrapper.output_indices,

len(forward_wrapper.output_tangents))

def _wrap_forward_function_with_jvps(

self, forward_function, backward_function,

inference_args, input_tangents):

"""Adds inline JVP computation to a forward function."""

forward_wrapper_graph = func_graph_module.FuncGraph(

_forward_name(self._func_graph.name))

with forward_wrapper_graph.as_default():

# Tell forward accumulators to free up space for new JVP computations,

# since one may be in the process of computing a JVP (if that computation

# triggered this function building).

#

# We'll make symbolic versions of input JVPs, run the forward function

# under forward accumulators to get symbolic output JVPs, then set those

# as outputs of the new wrapped forward function.

with forwardprop_util.push_forwardprop_state():

forward_captures = {

ops.tensor_id(internal): external

for external, internal in self._func_graph.captures}

for input_index, real_input in enumerate(self._func_graph.inputs):

# This loop is more or less equivalent to running tf.identity on each

# of self._func_graph.inputs. However, doing that also captures jvps

# for resource handles, which confuses the jvp capturing code below

# (since primal inputs are interwoven with jvp inputs).

input_placeholder = array_ops.placeholder(

dtype=real_input.dtype,

shape=real_input.shape)