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

if not data.dtype.is_ref_dtype:

return array_ops.identity(data, name=name)

else:

if not data.dtype.is_ref_dtype:

return next_iteration(data, name=name)

else:

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

return gen_control_flow_ops._ref_merge(values, name)

else:

use_ref=True, name=None):

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

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,

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

if data.dtype.is_ref_dtype:

return gen_control_flow_ops._ref_exit(data, name)

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

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

if isinstance(data, ops.Tensor):

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

else:

val, ind, dense_shape = data.values, data.indices, data.dense_shape

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

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

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

for inp in inputs]

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

return _Merge(inputs, name=name)

else:

inputs = math_ops._as_indexed_slices_list(inputs)

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

indices, chosen_index = _Merge(

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

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

if not all(inp.dense_shape for inp in inputs):

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

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

dense_shape, _ = _Merge(

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

else:

dense_shape = None

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

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

with ops.device(data.device):

if isinstance(data, ops.Tensor):

if not data.dtype.is_ref_dtype:

return switch(data, pred, name=name)

else:

return ref_switch(data, pred, name=name)

else:

"""A wrapper class for Operation.

class _ControlFlowOpInputs(object):

"""An indirection to capture the input tensors needed in backprop."""

def __init__(self, op, grad_state):

self._op = op

self._grad_state = grad_state

self._inputs = None

def __len__(self):

return len(self._op._inputs)

def __getitem__(self, index):

if self._inputs is None:

self._inputs = [None for _ in self._op.inputs]

if isinstance(index, int):

val = self._inputs[index]

if val is None:

f_val = self._op.inputs[index]

val = self._grad_state.GetRealValue(f_val)

self._inputs[index] = val

return val

elif isinstance(index, slice):

start, stop, step = index.indices(len(self))

vals = [self[i] for i in xrange(start, stop, step)]

return vals

else:

raise TypeError("index must be an integer or slice")

class _ControlFlowOpOutputs(object):

"""An indirection to capture the output tensors needed in backprop."""

def __init__(self, op, grad_state):

self._op = op

self._grad_state = grad_state

self._outputs = None

def __len__(self):

return len(self._op._outputs)

def __getitem__(self, index):

if self._outputs is None:

self._outputs = [None for _ in self._op.outputs]

if isinstance(index, int):

val = self._outputs[index]

if val is None:

f_val = self._op.outputs[index]

val = self._grad_state.GetRealValue(f_val)

self._outputs[index] = val

return val

elif isinstance(index, slice):

start, stop, step = index.indices(len(self))

vals = [self[i] for i in xrange(start, stop, step)]

return vals

else:

raise TypeError("index must be an integer or slice")

def __init__(self, op, grad_state):

self._grad_state = grad_state # The GradLoopState this op belongs to.

self._op = op

self._inputs = None

self._outputs = None

@property

def grad_state(self):

return self._grad_state

@property

def inputs(self):

if self._inputs is None:

self._inputs = self._ControlFlowOpInputs(self._op, self._grad_state)

return self._inputs

@property

def outputs(self):

if self._outputs is None:

self._outputs = self._ControlFlowOpOutputs(self._op, self._grad_state)

return self._outputs

@property

def op(self):

return self._op

@property

def name(self):

"""Returns the name of this instance of op."""

return self._op.name

@property

def _id(self):

"""Returns the unique id of this operation."""

return self._op._id

@property

def device(self):

"""Returns the device of this operation.

return self._op.device

@property

def type(self):

"""Returns the type of the op."""

return self._op.type

@property

def graph(self):

"""The `Graph` that contains this operation."""

return self._op.graph

def get_attr(self, name):

"""Returns the value of the attr of this op with the given `name`."""

return self._op.get_attr(name)

def _get_control_flow_context(self):

"""Returns the control flow context of this op."""

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

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

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.

# TODO(yuanbyu): Make sure the colocation of stack ops and value.

# pylint: disable=protected-access

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

# pylint: enable=protected-access

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

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)

# Protect stack push and order it before forward_index.

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

self.forward_context.Exit()

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)

value_ctxt.outer_context.Exit()

# Guard with a switch but take the other branch.

pred = self.history_map.get(value_ctxt.pred.name)

branch = value_ctxt.branch

value_ctxt.AddName(push.name)

value_ctxt.Enter()

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

value_ctxt.Exit()

else:

value_ctxt.Enter()

push = gen_data_flow_ops._stack_push(enter_acc, value)

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

if cond_ctxt:

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

grad_state = self

pred = None

while not pred 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)

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.

assert value.op.type != "Variable"

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

if real_value is None:

if _IsLoopConstantEnter(value.op):

# Special case for loop invariant.

if self._outer_grad_state:

# This is a nested loop so we record the history of this

# value in outer_forward_ctxt.

self._grad_context.Exit()

outer_value = value.op.inputs[0]

history_value = self._outer_grad_state.AddForwardAccumulator(

outer_value)

self._grad_context.Enter()

else:

# Just use the input value of this Enter node.

real_value = GetRealOp(value.op).inputs[0]

else:

# Record the history of this value in forward_ctxt.

# NOTE(yuanbyu): Don't record for constants.

self._grad_context.Exit()

history_value = self.AddForwardAccumulator(value)

self._grad_context.Enter()

if real_value is None:

# Add the stack pop op in the grad context.

real_value = self.AddBackPropAccumulatedValue(history_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):

forward_ctxt = _GetWhileContext(op)

if forward_ctxt is None:

return None

return self._map.get(forward_ctxt)

def MakeWrapper(self, op):

"""Make a wrapper for op if it is in a WhileContext."""

grad_state = self._GetGradState(op)

if grad_state:

return ControlFlowOpWrapper(op, grad_state)

return op

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

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

grad_state = self._GetGradState(op)

if grad_state:

grad_state.grad_context.Enter()

def ExitGradWhileContext(self, op):

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

grad_state = self._GetGradState(op)

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:

history_val = outer_grad_state.AddForwardAccumulator(val)

outer_grad_ctxt = outer_grad_state.grad_context

outer_grad_ctxt.Enter()

real_val = outer_grad_state.AddBackPropAccumulatedValue(

history_val, val)

result = array_ops.zeros_like(real_val)

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 = op.type in {"Switch", "RefSwitch"}

forward_ctxt = _GetWhileContext(op)

if forward_ctxt is None:

return array_ops.zeros_like(op.outputs[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.

op_ctxt.Enter()

zeros_shape = shape(val)

op_ctxt.Exit()

# Add forward accumulator for shape.

grad_state.grad_context.Exit()

history_shape = grad_state.AddForwardAccumulator(zeros_shape, dead_branch)

grad_state.grad_context.Enter()

# Create a zero tensor with the right shape.

shape = grad_state.AddBackPropAccumulatedValue(

history_shape, zero_shape, dead_branch)

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

"""Get the real op by removing the wrapper."""

while isinstance(op, ControlFlowOpWrapper):

op = op.op

"""Create the state for all the while loops involved in one gradients().

loop_state = None

for op in between_op_list:

if _IsLoopExit(op):

if loop_state is None:

loop_state = ControlFlowState()

loop_state.AddWhileContext(op, between_op_list, between_ops)

"""Return true if `op` is the Switch for a While loop."""

if op.type == "Switch" or op.type == "RefSwitch":

ctxt = op._get_control_flow_context()

return ctxt and isinstance(ctxt, WhileContext)

"""The base class for control flow context.

def __init__(self):

self._outer_context = ops.get_default_graph()._get_control_flow_context()

self._context_stack = []

# Values that have been already seen in this context.

self._values = set()

# Values referenced by but external to this context.

self._external_values = {}

@property

def outer_context(self):

"""Return the context containing this context."""

return self._outer_context

def AddName(self, name):

self._values.add(name)

# pylint: disable=protected-access

def Enter(self):

"""Enter this control flow context."""

graph = ops.get_default_graph()

self._context_stack.append(graph._get_control_flow_context())

graph._set_control_flow_context(self)

def Exit(self):

"""Exit this control flow context."""

graph = ops.get_default_graph()

last_context = self._context_stack.pop()

graph._set_control_flow_context(last_context)

def ExitResult(self, result):

"""Make a list of tensors available in the outer context."""

if self._outer_context:

for x in result:

self._outer_context.AddName(x.name)

def GetWhileContext(self):

"""Return the while context containing this context."""

if self._outer_context:

return self._outer_context.GetWhileContext()

return None

def MaybeAddToWhileContext(self, op):

"""Add a control dependency to the containing WhileContext.

while_ctxt = self.GetWhileContext()

if while_ctxt is not None:

# pylint: disable=protected-access

op._add_control_input(while_ctxt.GetControlPivot().op)