"""Return a tensor with the same shape and contents as the input tensor.

data = ops.convert_to_tensor_or_indexed_slices(data, as_ref=True)

if isinstance(data, ops.Tensor):

if data.dtype.is_ref_dtype:

return gen_array_ops._ref_identity(data, name=name)

else:

return array_ops.identity(data, name=name)

else:

values = _Identity(data.values, name=name)

indices = array_ops.identity(data.indices, name="indices")

if isinstance(data, ops.IndexedSlices):

dense_shape = data.dense_shape

if dense_shape is not None:

dense_shape = array_ops.identity(dense_shape, name="dense_shape")

return ops.IndexedSlices(values, indices, dense_shape)

elif isinstance(data, ops.SparseTensor):

dense_shape = array_ops.identity(data.shape, name="dense_shape")

return ops.SparseTensor(indices, values, dense_shape)

else:

data = ops.convert_to_tensor_or_indexed_slices(data, as_ref=True)

if isinstance(data, ops.Tensor):

if data.dtype.is_ref_dtype:

return ref_next_iteration(data, name=name)

else:

return next_iteration(data, name=name)

else:

values = _NextIteration(data.values, name=name)

indices = next_iteration(data.indices, name="indices")

if isinstance(data, ops.IndexedSlices):

dense_shape = data.dense_shape

if dense_shape is not None:

dense_shape = next_iteration(dense_shape, name="dense_shape")

return ops.IndexedSlices(values, indices, dense_shape)

elif isinstance(data, ops.SparseTensor):

dense_shape = next_iteration(data.shape, name="dense_shape")

return ops.SparseTensor(indices, values, dense_shape)

else:

use_ref=True, name=None):

"""Creates or finds a child frame, and makes `data` available to it.

data = ops.convert_to_tensor_or_indexed_slices(data, as_ref=True)

if isinstance(data, ops.Tensor):

if data.dtype.is_ref_dtype and use_ref:

return ref_enter(data, frame_name, is_constant, parallel_iterations,

name=name)

else:

return enter(data, frame_name, is_constant, parallel_iterations,

name=name)

else:

values = _Enter(data.values, frame_name, is_constant,

parallel_iterations, name=name)

indices = enter(data.indices, frame_name, is_constant,

parallel_iterations, name="indices")

if isinstance(data, ops.IndexedSlices):

dense_shape = data.dense_shape

if dense_shape is not None:

dense_shape = enter(dense_shape, frame_name, is_constant,

parallel_iterations, name="dense_shape")

return ops.IndexedSlices(values, indices, dense_shape)

elif isinstance(data, ops.SparseTensor):

dense_shape = enter(data.shape, frame_name, is_constant,

parallel_iterations, name="dense_shape")

return ops.SparseTensor(indices, values, dense_shape)

else:

"""Exits the current frame to its parent frame.

data = ops.convert_to_tensor_or_indexed_slices(data, as_ref=True)

if isinstance(data, ops.Tensor):

if data.dtype.is_ref_dtype:

return gen_control_flow_ops._ref_exit(data, name)

else:

return gen_control_flow_ops._exit(data, name)

else:

values = exit(data.values, name=name)

indices = gen_control_flow_ops._exit(data.indices, name="indices")

if isinstance(data, ops.IndexedSlices):

dense_shape = data.dense_shape

if dense_shape is not None:

dense_shape = gen_control_flow_ops._exit(dense_shape, name)

return ops.IndexedSlices(values, indices, dense_shape)

elif isinstance(data, ops.SparseTensor):

dense_shape = gen_control_flow_ops._exit(data.shape, name)

return ops.SparseTensor(indices, values, dense_shape)

else:

"""Forwards `data` to an output determined by `pred`.

with ops.op_scope([data, pred], name, "Switch") as name:

data = ops.convert_to_tensor_or_indexed_slices(data, dtype=dtype,

name="data", as_ref=True)

pred = ops.convert_to_tensor(pred, name="pred")

if isinstance(data, ops.Tensor):

return gen_control_flow_ops._switch(data, pred, name=name)

else:

if not isinstance(data, (ops.IndexedSlices, ops.SparseTensor)):

raise TypeError("Type %s not supported" % type(data))

val, ind = data.values, data.indices

val_f, val_t = gen_control_flow_ops._switch(val, pred, name=name)

ind_f, ind_t = gen_control_flow_ops._switch(ind, pred, name="indices")

if isinstance(data, ops.IndexedSlices):

dense_shape = data.dense_shape

if dense_shape is not None:

dense_shape_f, dense_shape_t = gen_control_flow_ops._switch(

dense_shape, pred, name="dense_shape")

else:

dense_shape_f, dense_shape_t = None, None

return (ops.IndexedSlices(val_f, ind_f, dense_shape_f),

ops.IndexedSlices(val_t, ind_t, dense_shape_t))

else:

dense_shape = data.shape

dense_shape_f, dense_shape_t = gen_control_flow_ops._switch(

data.shape, pred, name="dense_shape")

return (ops.SparseTensor(ind_f, val_f, dense_shape_f),

"""Forwards `data` to an output determined by `pred`.

data = ops.convert_to_tensor_or_indexed_slices(data, name="data")

# NOTE(vrv): ops.colocate_with(data, ignore_existing=True) below

# addresses the following scenario.

#

# Assume you execute Optimizer.apply_gradients() in a branch of a cond().

#

# 1. The update op is created inside a `with ops.colocate(var):` block

#

# 2. Some tensor `data` is captured and a switch is created in a

# `with ops.colocate_with(data):` block.

#

# with ops.colocate_with(var):

# with ops.colocate_with(data):

# op = ...

#

# var and data may be pinned to different devices, so we want to ops

# created within ops.colocate_with(data) to ignore the existing stack.

with ops.colocate_with(data, ignore_existing=True):

if isinstance(data, ops.Tensor):

if data.dtype.is_ref_dtype:

return ref_switch(data, pred, name=name)

"""Returns the value of an available element of `inputs`.

with ops.op_scope(inputs, name, "Merge") as name:

inputs = [ops.convert_to_tensor_or_indexed_slices(inp, as_ref=True)

for inp in inputs]

if all([isinstance(v, ops.Tensor) for v in inputs]):

if all([v.dtype.is_ref_dtype for v in inputs]):

return gen_control_flow_ops._ref_merge(inputs, name)

else:

return gen_control_flow_ops._merge(inputs, name)

elif all([isinstance(v, ops.SparseTensor) for v in inputs]):

# Only handle the case when all inputs are SparseTensor.

values, _ = merge([inp.values for inp in inputs], name=name)

indices, chosen_index = gen_control_flow_ops._merge(

[inp.indices for inp in inputs], name="indices")

dense_shape, _ = gen_control_flow_ops._merge(

[inp.shape for inp in inputs], name="dense_shape")

return ops.SparseTensor(indices, values, dense_shape), chosen_index

else:

# For now convert all the inputs as IndexedSlices.

inputs = math_ops._as_indexed_slices_list(inputs)

values, _ = merge([inp.values for inp in inputs], name=name)

indices, chosen_index = gen_control_flow_ops._merge(

[inp.indices for inp in inputs], name="indices")

if any(inp.dense_shape is not None for inp in inputs):

if any(inp.dense_shape is None for inp in inputs):

raise ValueError("Either all merged IndexedSlices must have a "

"dense_shape, or none must have a dense_shape.")

dense_shape, _ = gen_control_flow_ops._merge(

[inp.dense_shape for inp in inputs], name="dense_shape")

else:

dense_shape = None

new_ta = tensor_array_ops.TensorArray(

dtype=ta.dtype, handle=ta.handle, flow=t_or_flow,

infer_shape=ta._infer_shape)

new_ta._elem_shape = ta._elem_shape

if len(tensors_or_tensorarrays) != len(tensors_or_flows):

raise ValueError(

"Lengths of original Tensor list and new list do not match: %d vs. %d"

% (len(tensors_or_tensorarrays), len(tensors_or_flows)))

return [

_make_tensor_array(ta, t_or_flow)

if isinstance(ta, tensor_array_ops.TensorArray)

else t_or_flow

"""Return the enter op if we can infer `value` to be a loop invariant."""

id_ops = {"Switch", "RefSwitch", "Identity", "RefIdentity"}

op = value.op

while op.type in id_ops:

op = op.inputs[0].op

if shape1.dims is None or shape1.ndims != shape2.ndims:

return tensor_shape.unknown_shape()

rdims = []

for dim1, dim2 in zip(shape1.dims, shape2.dims):

if dim1 == dim2:

rdims.append(dim1)

else:

rdims.append(tensor_shape.Dimension(None))

"""Add NextIteration and back edge from v to m."""

if isinstance(m, ops.Tensor):

v = ops.convert_to_tensor(v)

v = _NextIteration(v)

m.op._update_input(1, v) # pylint: disable=protected-access

elif isinstance(m, ops.IndexedSlices):

# pylint: disable=protected-access

v = math_ops._as_indexed_slices(v)

v = _NextIteration(v)

m.values.op._update_input(1, v.values)

m.indices.op._update_input(1, v.indices)

# pylint: enable=protected-access

if m.dense_shape is not None:

if v.dense_shape is None:

raise ValueError("Must have dense shape: %s" % v.name)

m.dense_shape.op._update_input(1, v.dense_shape)

elif isinstance(m, ops.SparseTensor):

if not isinstance(v, ops.SparseTensor):

raise ValueError("Must be a sparse tensor: %s" % v.name)

v = _NextIteration(v)

# pylint: disable=protected-access

m.values.op._update_input(1, v.values)

m.indices.op._update_input(1, v.indices)

m.shape.op._update_input(1, v.shape)

# pylint: enable=protected-access

else:

raise TypeError("Type %s not supported" % type(m))

"""The state used for constructing the gradient graph for a while loop.

def __init__(self, forward_ctxt, outer_grad_state):

# The grad loop state for the outer while loop.

self._outer_grad_state = None

# The while loop context for forward.

self._forward_context = None

# The loop counter added by AddForwardCounter. It is the value

# of the loop counter for the next iteration.

self._forward_index = None

# A sync op for forward.

self._forward_sync = None

# The while loop context for backprop.

self._grad_context = None

# The loop counter added by AddBackPropCounter. It is the value

# of the loop counter for the current iteration.

self._grad_index = None

# A sync op for backprop.

self._grad_sync = None

# Information needed by backprop.

self._history_map = {}

self._switch_map = {}

self._outer_grad_state = outer_grad_state

if outer_grad_state:

outer_forward_ctxt = outer_grad_state.forward_context

else:

outer_forward_ctxt = forward_ctxt.outer_context

# Add the forward loop counter.

if outer_forward_ctxt: outer_forward_ctxt.Enter()

cnt, forward_index = forward_ctxt.AddForwardCounter()

if outer_forward_ctxt: outer_forward_ctxt.Exit()

self._forward_context = forward_ctxt

self._forward_index = forward_index

# Add the backprop WhileContext, and the backprop loop counter.

if outer_grad_state:

# This is a nested loop. Remember the iteration counts for each

# execution of this inner loop.

outer_forward_ctxt.AddName(cnt.name)

history_cnt = outer_grad_state.AddForwardAccumulator(cnt)

outer_grad_ctxt = outer_grad_state.grad_context

outer_grad_ctxt.Enter()

self._grad_context = WhileContext(forward_ctxt.parallel_iterations,

forward_ctxt.back_prop,

forward_ctxt.swap_memory,

forward_ctxt.name,

self)

real_cnt = outer_grad_state.AddBackPropAccumulatedValue(history_cnt, cnt)

self._grad_index = self._grad_context.AddBackPropCounter(real_cnt)

outer_grad_ctxt.Exit()

else:

if outer_forward_ctxt: outer_forward_ctxt.Enter()

self._grad_context = WhileContext(forward_ctxt.parallel_iterations,

forward_ctxt.back_prop,

forward_ctxt.swap_memory,

forward_ctxt.name,

self)

self._grad_index = self._grad_context.AddBackPropCounter(cnt)

if outer_forward_ctxt: outer_forward_ctxt.Exit()

@property

def outer_grad_state(self):

"""The grad loop state for outer loop."""

return self._outer_grad_state

@property

def forward_context(self):

"""The while loop context for forward."""

return self._forward_context

@property

def forward_index(self):

"""The loop index of forward loop."""

return self._forward_index

@property

def forward_sync(self):

"""A control trigger node for synchronization in the forward loop.

if self._forward_sync is None:

with ops.control_dependencies(None):

self._forward_sync = control_trigger(name="f_sync")

self._forward_sync._set_control_flow_context(self._forward_context)

self._forward_index.op._add_control_input(self._forward_sync)

return self._forward_sync

@property

def grad_context(self):

"""The corresponding WhileContext for gradient."""

return self._grad_context

@property

def grad_index(self):

"""The loop index of backprop loop."""

return self._grad_index

@property

def grad_sync(self):

"""A control trigger node for synchronization in the grad loop.

if self._grad_sync is None:

with ops.control_dependencies(None):

self._grad_sync = control_trigger(name="b_sync")

self._grad_sync._set_control_flow_context(self._grad_context)

self._grad_index.op._add_control_input(self._grad_sync)

return self._grad_sync

@property

def history_map(self):

"""The map that records all the tensors needed for backprop."""

return self._history_map

@property

def switch_map(self):

"""The map that records all the Switch ops for the While loop."""

return self._switch_map

def AddForwardAccumulator(self, value, dead_branch=False):

"""Add an accumulator for each forward tensor that is needed in backprop.

curr_ctxt = ops.get_default_graph()._get_control_flow_context()

with ops.control_dependencies(None):

if curr_ctxt: curr_ctxt.Enter()

with ops.colocate_with(value):

# pylint: disable=protected-access

acc = gen_data_flow_ops._stack(value.dtype.base_dtype, name="f_acc")

# pylint: enable=protected-access

if curr_ctxt: curr_ctxt.Exit()

# Make acc available in the forward context.

enter_acc = self.forward_context.AddValue(acc)

# Add the stack_push op in the context of value.op.

swap_enabled = self.forward_context.swap_memory

value_ctxt = value.op._get_control_flow_context()

if _IsLoopExit(value.op):

value_ctxt = value_ctxt.outer_context

if value_ctxt == self.forward_context:

# value is not nested in the forward context.

self.forward_context.Enter()

push = gen_data_flow_ops._stack_push(

enter_acc, value, swap_memory=swap_enabled)

self.forward_context.Exit()

# Protect stack push and order it before forward_index.

self.forward_index.op._add_control_input(push.op)

else:

# value is in a cond context within the forward context.

assert isinstance(value_ctxt, CondContext)

if dead_branch:

# The special case for creating a zero tensor for a dead

# branch of a switch. See ControlFlowState.ZerosLike().

value_ctxt.outer_context.Enter()

push = gen_data_flow_ops._stack_push(

enter_acc, value, swap_memory=swap_enabled)

value_ctxt.outer_context.Exit()

push.op._set_control_flow_context(value_ctxt)

else:

value_ctxt.Enter()

push = gen_data_flow_ops._stack_push(

enter_acc, value, swap_memory=swap_enabled)

value_ctxt.Exit()

# Protect stack push and order it before forward_sync.

self.forward_sync._add_control_input(push.op)

# Order stack push after the successor of forward_index

add_op = self.forward_index.op.inputs[0].op

push.op._add_control_input(add_op)

return acc

def AddBackPropAccumulatedValue(self, history_value, value,

dead_branch=False):

"""Add the getter for an accumulated value in the grad context.

history_ctxt = history_value.op._get_control_flow_context()

# Find the cond context that controls history_value.

cond_ctxt = None

value_ctxt = value.op._get_control_flow_context()

while value_ctxt and value_ctxt != history_ctxt:

if isinstance(value_ctxt, CondContext):

cond_ctxt = value_ctxt

break

value_ctxt = value_ctxt.outer_context

with ops.control_dependencies(None):

self.grad_context.Enter()

if cond_ctxt:

# Guard stack pop with a switch if it is controlled by a cond

grad_state = self

pred = None

while pred is None and grad_state:

pred = grad_state.history_map.get(cond_ctxt.pred.name)

grad_state = grad_state.outer_grad_state

branch = (1 - cond_ctxt.branch) if dead_branch else cond_ctxt.branch

history_value = _SwitchRefOrTensor(history_value, pred)[branch]

pop = gen_data_flow_ops._stack_pop(history_value, value.dtype.base_dtype)

self.grad_context.Exit()

if self.grad_context.parallel_iterations > 1:

# All pops are ordered after pivot_for_body and before grad_sync.

self.grad_sync._add_control_input(pop.op)

return pop

def GetRealValue(self, value):

"""Get the real value of `value`.

assert value.op.type != "Variable"

real_value = self._history_map.get(value.name)

if real_value is None:

cur_value = value

cur_grad_state = self

while True:

enter_op = _GetLoopConstantEnter(cur_value)

if enter_op:

# Special case: cur_value comes from a constant Enter node.

cur_value = enter_op.inputs[0]

cur_grad_state = cur_grad_state.outer_grad_state

if cur_grad_state is None:

# We are now outside all nested loops for this gradient(),

# so `value` is a loop invariant and there is no need to

# save the history of value. Just make cur_value to enter

# the right control flow context.

real_value = self._grad_context.AddValue(cur_value)

break

else:

# Record the history of this value in forward_ctxt.

# TODO(yuanbyu): Avoid recording constants.

self._grad_context.Exit()

h_value = cur_grad_state.AddForwardAccumulator(cur_value)

self._grad_context.Enter()

break

if real_value is None:

# Add the stack pop op in the grad context.

real_value = self.AddBackPropAccumulatedValue(h_value, value)

self._history_map[value.name] = real_value

"""Get the WhileContext to which this op belongs."""

ctxt = op._get_control_flow_context()

if ctxt:

ctxt = ctxt.GetWhileContext()

"""Maintain the mapping from the loops to their grad states."""

def __init__(self):

self._map = {} # maps forward loop context to GradLoopState

def _GetGradState(self, op, before):

"""Return the grad state for this op if it's in a forward loop context."""

if before and _IsLoopExit(op):

forward_ctxt = op._get_control_flow_context()

forward_ctxt = forward_ctxt.outer_context

if forward_ctxt:

forward_ctxt = forward_ctxt.GetWhileContext()

else:

forward_ctxt = _GetWhileContext(op)

if forward_ctxt:

return self._map.get(forward_ctxt)

return None

def GetAllLoopExits(self):

"""Return a list containing the exits of all the loops."""

loop_exits = []

for forward_ctxt in self._map:

for loop_exit in forward_ctxt.loop_exits:

loop_exits.append(loop_exit)

return loop_exits

def EnterGradWhileContext(self, op, before):

"""Enter the WhileContext for gradient computation."""

grad_state = self._GetGradState(op, before)

if grad_state:

grad_state.grad_context.Enter()

def ExitGradWhileContext(self, op, before):

"""Exit the WhileContext for gradient computation."""

grad_state = self._GetGradState(op, before)

if grad_state:

grad_state.grad_context.Exit()

def AddWhileContext(self, op, between_op_list, between_ops):

"""Add the grad state for the while loop that op belongs to.

forward_ctxt = _GetWhileContext(op)

grad_state = self._map.get(forward_ctxt)

if grad_state is None:

# This is a new while loop so create a grad state for it.

outer_forward_ctxt = forward_ctxt.outer_context

if outer_forward_ctxt:

outer_forward_ctxt = outer_forward_ctxt.GetWhileContext()

outer_grad_state = None

if outer_forward_ctxt:

outer_grad_state = self._map.get(outer_forward_ctxt)

grad_state = GradLoopState(forward_ctxt, outer_grad_state)

self._map[forward_ctxt] = grad_state

# We need to include all exits of a loop for backprop.

for loop_exit in forward_ctxt.loop_exits:

if not between_ops[loop_exit.op._id]:

between_ops[loop_exit.op._id] = True

between_op_list.append(loop_exit.op)

def ZerosLikeForExit(self, val):

"""Create zeros_like gradient for a loop exit.

val_shape = val.get_shape()

forward_ctxt = val.op._get_control_flow_context()

outer_forward_ctxt = forward_ctxt.outer_context

if outer_forward_ctxt:

outer_forward_ctxt = outer_forward_ctxt.GetWhileContext()

outer_grad_state = None

if outer_forward_ctxt:

outer_grad_state = self._map.get(outer_forward_ctxt)

if outer_grad_state:

# This is a nested loop.

if val_shape.is_fully_defined():

# If the shape is known statically, just create a zero tensor

# with the right shape in the right context.

outer_grad_state.grad_context.Enter()

result = array_ops.zeros(val_shape.dims, val.dtype)

outer_grad_state.grad_context.Exit()

else:

# Only the shape of value is needed for backprop.

forward_ctxt.outer_context.Enter()

shape = array_ops.shape(value)

forward_ctxt.outer_context.Exit()

# Save the shape to a stack.

history_shape = outer_grad_state.AddForwardAccumulator(shape)

# Get the shape back from the stack.

outer_grad_ctxt = outer_grad_state.grad_context

outer_grad_ctxt.Enter()

real_shape = outer_grad_state.AddBackPropAccumulatedValue(

history_shape, shape)

result = array_ops.zeros(real_shape, value.dtype)

outer_grad_ctxt.Exit()

else:

# This is not a nested loop.

if val_shape.is_fully_defined():

# If the shape is known statically, just create a zero tensor

# with the right shape.

result = array_ops.zeros(val_shape.dims, val.dtype)

else:

result = array_ops.zeros_like(val)

return result

def ZerosLike(self, op, index):

"""Create zeros_like for the specified output of an op.

if IsLoopSwitch(op): return None

dead_branch = IsSwitch(op)

forward_ctxt = _GetWhileContext(op)

if forward_ctxt is None:

# op is not in a while loop that is part of gradients().

return ZerosLikeOutsideLoop(op, index)

op_ctxt = op._get_control_flow_context()

grad_state = self._map.get(forward_ctxt)

val = ops.convert_to_tensor(op.outputs[index], name="tensor")

shape = val.get_shape()

if shape.is_fully_defined():

# If the shape is known statically, just create a zero tensor with

# the right shape in the grad loop context.

result = constant_op.constant(0, shape=shape.dims, dtype=val.dtype)

if dead_branch:

# op is a cond switch. Guard the zero tensor with a switch.

pred = grad_state.history_map.get(op_ctxt.pred.name)

branch = op_ctxt.branch

result = _SwitchRefOrTensor(result, pred)[1 - branch]

else:

# Unknown shape so keep a history of the shape at runtime.

if dead_branch:

# Need to add a special switch to guard the value.

pred = op_ctxt.pred

branch = op_ctxt.branch

op_ctxt.outer_context.Enter()

val = _SwitchRefOrTensor(op.inputs[0], pred)[1 - branch]

zeros_shape = array_ops.shape(val)

op_ctxt.outer_context.Exit()

val.op._set_control_flow_context(op_ctxt)

zeros_shape.op._set_control_flow_context(op_ctxt)

else:

op_ctxt.Enter()

zeros_shape = array_ops.shape(val)

op_ctxt.Exit()

# Add forward accumulator for shape.

grad_state.grad_context.Exit()

h_shape = grad_state.AddForwardAccumulator(

zeros_shape, dead_branch=dead_branch)

grad_state.grad_context.Enter()

# Create a zero tensor with the right shape.

shape = grad_state.AddBackPropAccumulatedValue(

h_shape, zeros_shape, dead_branch)

result = array_ops.zeros(shape, val.dtype)

return result

def PostProcessing(self):

"""Perform postprocessing at the end of gradients().

for _, grad_state in self._map.items():

for _, b_merge in grad_state.switch_map.items():

if b_merge.op.inputs[0] == b_merge.op.inputs[1]:

# The value of this loop variable at iteration i+1 doesn't

# depend on its value at iteration i. So use zeros as the

# gradients for all iterations > 0.

dtype = b_merge.op.inputs[0].dtype

shape = b_merge.op.inputs[0].get_shape()

if not shape.is_fully_defined():

shape = None

grad_state.grad_context.Enter()

grad_val = constant_op.constant(0, dtype=dtype, shape=shape)

grad_val = _NextIteration(grad_val)

grad_state.grad_context.Exit()

# pylint: disable=protected-access

if not shape: