path: root/tensorflow/python/ops/control_flow_ops.py

"""## Control Flow Operations

TensorFlow provides several operations and classes that you can use to control
the execution of operations and add conditional dependencies to your graph.

@@identity
@@tuple
@@group
@@no_op
@@count_up_to

## Logical Operators

TensorFlow provides several operations that you can use to add logical operators
to your graph.

@@logical_and
@@logical_not
@@logical_or
@@logical_xor

## Comparison Operators

TensorFlow provides several operations that you can use to add comparison
operators to your graph.

@@equal
@@not_equal
@@less
@@less_equal
@@greater
@@greater_equal
@@select
@@where

## Debugging Operations

TensorFlow provides several operations that you can use to validate values and
debug your graph.

@@is_finite
@@is_inf
@@is_nan
@@verify_tensor_all_finite
@@check_numerics
@@add_check_numerics_ops
@@Assert
@@Print
"""
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import types
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import common_shapes
from tensorflow.python.ops import constant_op
from tensorflow.python.ops import gen_control_flow_ops
from tensorflow.python.ops import gen_array_ops
from tensorflow.python.ops import logging_ops
from tensorflow.python.ops import math_ops
# pylint: disable=wildcard-import,undefined-variable
from tensorflow.python.ops.gen_control_flow_ops import *


# We override the 'tuple' for a control flow op, so we keep python's
# existing 'tuple' for later use in this module.
_basetuple = tuple


# pylint: disable=protected-access
def _Identity(data, name=None):
  """Return a tensor with the same shape and contents as the input tensor.

  Args:
    data: A Tensor.
    name: A name for this operation (optional).

  Returns:
    A Tensor with the same type and value as the input Tensor.
  """
  if not data.dtype.is_ref_dtype:
    return array_ops.identity(data, name=name)
  else:
    return gen_array_ops._ref_identity(data, name=name)


def _Enter(data, frame_name, is_constant=False, parallel_iterations=10,
           name=None):
  """Creates or finds a child frame, and makes 'data' available to it.

  The unique `frame_name` is used by the `Executor` to identify frames. If
  `is_constant` is true, `output` is a constant in the child frame; otherwise
  it may be changed in the child frame. At most `parallel_iterations` iterations
  are run in parallel in the child frame.

  Args:
    data: The tensor to be made available to the child frame.
    frame_name: The name of the child frame.
    is_constant: If true, the output is constant within the child frame.
    parallel_iterations: The number of iterations allowed to run in parallel.
    name: A name for this operation (optional).

  Returns:
    The same tensor as 'data'.
  """
  if not data.dtype.is_ref_dtype:
    return enter(data, frame_name, is_constant, parallel_iterations,
                 name=name)
  else:
    return ref_enter(data, frame_name, is_constant, parallel_iterations,
                     name=name)


def exit(data, name=None):
  """Exits the current frame to its parent frame.

  Exit makes its input `data` available to the parent frame.

  Args:
    data: The tensor to be made available to the parent frame.
    name: A name for this operation (optional).

  Returns:
    The same tensor as `data`.
  """
  return gen_control_flow_ops._exit(data, name)


def switch(data, pred, name=None):
  """Forwards `data` to an output determined by `pred`.

  If `pred` is true, the `data` input is forwared to the first output.
  Otherwise, the data goes to the second output.

  This op handles `Tensor`s and `IndexedSlices`.

  Args:
    data: The tensor to be forwarded to the appropriate output.
    pred: A scalar that specifies which output port will receive data.
    name: A name for this operation (optional).

  Returns:
    `(output_true, output_false)`: If `pred` is true, data will be forwarded to
    `output_true`, otherwise it goes to `output_false`.
  """
  with ops.op_scope([data, pred], name, "Switch") as name:
    data = ops.convert_to_tensor_or_indexed_slices(data, 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),
              ops.IndexedSlices(val_t, ind_t, dense_shape_t))


def merge(inputs, name=None):
  """Returns the value of an available element of `inputs`.

  This op tests each of the tensors in `inputs` in turn to determine if any of
  them is available. If it finds an available tensor, it returns it and its
  index in `inputs`.

  It is an error if more than one tensor in `inputs` is available. If no tensor
  in `inputs` is available, the returned tensor and index are not set.

  This op handles both `Tensor`s and `IndexedSlices`. If inputs has a mix of
  `Tensor`s and `IndexedSlices`, all inputs are converted to IndexedSlices
  before merging.

  Args:
    inputs: The input tensors, at most one of which is available.
    name: A name for this operation (optional).

  Returns:
    A tuple containing the chosen input tensor and its index in `inputs`.

  Raises:
    ValueError: If inputs are IndexedSlices and some but not all have a
      dense_shape property.
  """
  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 gen_control_flow_ops._merge(inputs, name=name)
    else:
      inputs = math_ops._as_indexed_slices_list(inputs)
      values, _ = gen_control_flow_ops._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 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, _ = gen_control_flow_ops._merge(
            [inp.dense_shape for inp in inputs], name="dense_shape")
      else:
        dense_shape = None
      return ops.IndexedSlices(values, indices, dense_shape), chosen_index


def _SwitchRefOrTensor(data, pred, name="Switch"):
  """Forwards `data` to an output determined by `pred`.

  If `pred` is true, the `data` input is forwared to the first output.
  Otherwise, the data goes to the second output.

  This op handles `Tensor`s and `IndexedSlices`.

  Args:
    data: The tensor to be forwarded to the appropriate output.
    pred: A scalar that specifies which output port will receive data.
    name: A name for this operation (optional).

  Returns:
    `(output_false, output_false)`: If `pred` is true, data will be forwarded to
    `output_true`, otherwise it goes to `output_false`.

  Raises:
    TypeError: if data is not a Tensor or IndexedSlices
  """
  data = ops.convert_to_tensor_or_indexed_slices(data, name="data")
  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:
    return switch(data, pred, name=name)


class ControlFlowOpInputs(object):
  """An indirection to capture the input tensors needed in backprop."""

  def __init__(self, op):
    self._op = op
    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 = _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):
    self._op = op
    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 = _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")


class ControlFlowOpWrapper(object):
  """A wrapper class for Operation."""

  def __init__(self, op):
    self._op = op
    self._inputs = None
    self._outputs = None

  @property
  def inputs(self):
    if self._inputs is None:
      self._inputs = ControlFlowOpInputs(self._op)
    return self._inputs

  @property
  def outputs(self):
    if self._outputs is None:
      self._outputs = ControlFlowOpOutputs(self._op)
    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.

    Returns:
      a string or None if the device was not set.
    """
    return self._op.device

  @property
  def output_types(self):
    return self._op.output_types

  @property
  def input_types(self):
    return self._op._input_types

  @property
  def type(self):
    """Returns the type of the op."""
    return self._op.type

  @property
  def graph(self):
    """Returns the parent graph."""
    return self._op.graph

  def GetAttr(self, attr_name):
    """Returns the value of attribute 'attr_name' of NodeDef."""
    return self._op.get_attr(attr_name)

  def _get_control_flow_context(self):
    return self._op._get_control_flow_context()


def GetRealOp(op):
  while isinstance(op, ControlFlowOpWrapper):
    op = op.op
  return op


def MakeWrapper(op):
  """Make a wrapper for op if it is in a WhileContext."""
  forward_ctxt = op._get_control_flow_context()
  if forward_ctxt and isinstance(forward_ctxt, WhileContext):
    return ControlFlowOpWrapper(op)
  return op


def EnterGradWhileContext(op):
  """Enter the WhileContext for gradient computation."""
  forward_ctxt = op._get_control_flow_context()
  if forward_ctxt and isinstance(forward_ctxt, WhileContext):
    grad_ctxt = forward_ctxt.CreateGradWhileContext()
    grad_ctxt.Enter()


def ExitGradWhileContext(op):
  """Exit the WhileContext for gradient computation."""
  forward_ctxt = op._get_control_flow_context()
  if forward_ctxt and isinstance(forward_ctxt, WhileContext):
    assert forward_ctxt.grad_context
    forward_ctxt.grad_context.Exit()


def _GetRealValue(value):
  """Get the real value.

  If backprop "uses" a value produced by forward inference, an
  accumulator is added in the forward loop to accumulate its values,
  so we use the accumulated value, indexed by the backprop counter.

  Args:
    value: A tensor to be captured.

  Returns:
    The same tensor value from the saved history.
  """
  real_value = value
  forward_ctxt = value.op._get_control_flow_context()
  real_value = forward_ctxt.history_map.get(value.name)
  assert value.op.type != "Variable"
  if real_value is None:
    if value.op.type == "Enter" and value.op.get_attr("is_constant"):
      # Use the input of this Enter node
      real_value = GetRealOp(value.op).inputs[0]
    else:
      # Accumulate the history of this value.
      # NOTE(yuanbyu): Don't accumulate for constants. One approach is
      # to deepcopy the constants for the grad while context.
      history_value = forward_ctxt.AddForwardAccumulateLoop(value)

      # The shapes of the whole history and a single event element.
      forward_ctxt.grad_context.Exit()
      elem_rank = array_ops.rank(history_value) - 1
      elem_rank_vec = array_ops.expand_dims(elem_rank, 0)
      elem_shape = array_ops.slice(array_ops.shape(history_value), [1],
                                   elem_rank_vec)
      slice_shape = array_ops.concat(0, [[1], elem_shape])
      forward_ctxt.grad_context.Enter()

      # The begin position of the slice at slice_index.
      slice_index = forward_ctxt.grad_context.index
      b1 = array_ops.zeros(elem_rank_vec, dtype=types.int32)
      b = array_ops.concat(0, [array_ops.expand_dims(slice_index, 0), b1])

      # The slice at slice_index.
      # TODO(irving): Replace with gather once that's GPU accelerated
      real_value = array_ops.squeeze(
          array_ops.slice(history_value,
                          b,
                          slice_shape,
                          name="real"),
          squeeze_dims=[0])
  forward_ctxt.history_map[value.name] = real_value
  return real_value


def IsLoopSwitch(op):
  """Returns true if `op` is the Switch for a While loop."""
  if op.type == "Switch":
    ctxt = op._get_control_flow_context()
    return ctxt and isinstance(ctxt, WhileContext)
  return False


class ControlFlowContext(object):
  """The base class for control flow context.

  The usage pattern is a sequence of (Enter, Exit) followed by a final
  ExitResult.
  """

  def AddName(self, name):
    self._values.add(name)

  # pylint: disable=protected-access
  def Enter(self):
    """Enter the current context."""
    self._outer_context = ops.get_default_graph()._get_control_flow_context()
    ops.get_default_graph()._set_control_flow_context(self)

  def Exit(self):
    """Exit the current context."""
    ops.get_default_graph()._set_control_flow_context(self._outer_context)
  # pylint: enable=protected-access

  def ExitResult(self, result):
    """Make a list of tensors available in the outer context."""
    if self._outer_context is not None:
      for x in result:
        self._outer_context.AddName(x.name)

  def GetWhileContext(self):
    """Get the current while context."""
    if self._outer_context is not None:
      return self._outer_context.GetWhileContext()
    return None

  def AddToWhileContext(self, op):
    """Add a control dependency to the containing WhileContext.

    The added control dependency ensures that the outputs of this op
    belong to the WhileContext.

    Args:
      op: An operation.
    """
    while_ctxt = self.GetWhileContext()
    if while_ctxt is not None:
      # pylint: disable=protected-access
      op._add_control_input(while_ctxt.GetControlPivot().op)
      # pylint: enable=protected-access


class CondContext(ControlFlowContext):
  """The context for the conditional construct."""

  def __init__(self, pred, pivot, branch):
    self._pred = pred
    self._outer_context = None
    self._pivot = pivot
    self._branch = branch
    self._values = set()
    self._values.add(pred.name)
    self._values.add(pivot.name)
    self._external_values = {}

  @property
  def pred(self):
    return self._pred

  @property
  def pivot(self):
    return self._pivot

  @property
  def branch(self):
    return self._branch

  def AddValue(self, val):
    """Add 'val' to the current context and its outer context recursively."""
    result = val
    if val.name not in self._values:
      self._values.add(val.name)
      if self._outer_context is not None:
        result = self._outer_context.AddValue(val)
      result = with_dependencies([self._pivot], result)
      self._external_values[val.name] = result
    return result

  def AddOp(self, op):
    """Add 'op' to the current context."""
    if not op.inputs:
      # Add this op to the enclosing while context
      self.AddToWhileContext(op)
      # pylint: disable=protected-access
      op._add_control_input(self._pivot.op)
      # pylint: enable=protected-access
      for x in op.outputs:
        self._values.add(x.name)
    else:
      for index in range(len(op.inputs)):
        x = op.inputs[index]
        if x.name not in self._values:
          self._values.add(x.name)
          # Add this value to the parent contexts up to the context that
          # creates this value.
          real_x = x
          if self._outer_context is not None:
            real_x = self._outer_context.AddValue(x)
          real_x = _SwitchRefOrTensor(real_x, self._pred)[self._branch]
          self._external_values[x.name] = real_x
        x = self._external_values.get(x.name)
        if x is not None:
          op._update_input(index, x)
      for x in op.outputs:
        self._values.add(x.name)

  def BuildCondBranch(self, fn):
    """Add the subgraph defined by fn() to the graph."""
    r = fn()
    result = []
    if r is not None:
      if not isinstance(r, list) and not isinstance(r, _basetuple):
        r = [r]
      for v in r:
        if isinstance(v, ops.Operation):
          v = with_dependencies([v], self._pivot)
        elif v.name not in self._values:
          self._values.add(v.name)
          if self._outer_context is not None:
            v = self._outer_context.AddValue(v)
          v = _SwitchRefOrTensor(v, self._pred)[self._branch]
        else:
          external_v = self._external_values.get(v.name)
          if external_v is not None:
            v = external_v
        result.append(v)
    return result


def cond(pred, fn1, fn2, name=None):
  """Return either 'fn1()' or 'fn2()' based on the boolean predicate 'pred'.

  `fn1` and `fn2` both return lists of output tensors. `fn1` and `fn2` must have
  the same number and type of outputs.

  Args:
    pred: A scalar determining whether to return the result of `fn1` or `fn2`.
    fn1: The function to be performed if pred is true.
    fn2: The function to be performed if pref is false.
    name: Optional name prefix for the returned tensors.

  Returns:
    Tensors returned by the call to either `fn1` or `fn2`. If the functions
    return a singleton list, the element is extracted from the list.

  Raises:
    TypeError: if `fn1` or `fn2` is not callable.
    ValueError: if `fn1` and `fn2` do not return the same number of tensors, or
                return tensors of different types.

  Example:
  ```python
    x = constant(2)
    y = constant(5)
    def f1(): return constant(17)
    def f2(): return constant(23)
    r = cond(math_ops.less(x, y), f1, f2)
    # r is set to f1()
  ```
  """
  with ops.op_scope([pred], name, "Cond") as name:
    if not callable(fn1):
      raise TypeError("fn1 must be callable.")
    if not callable(fn2):
      raise TypeError("fn2 must be callable.")

    # Add the Switch to the graph.
    p_2, p_1 = switch(pred, pred)
    pivot_1 = array_ops.identity(p_1, name="switch_t")
    pivot_2 = array_ops.identity(p_2, name="switch_f")
    pred = array_ops.identity(pred, name="pred_id")

    # Build the graph for the true branch in a new context.
    context_t = CondContext(pred, pivot_1, 1)
    context_t.Enter()
    res_t = context_t.BuildCondBranch(fn1)
    context_t.ExitResult(res_t)
    context_t.Exit()

    # Build the graph for the false branch in a new context.
    context_f = CondContext(pred, pivot_2, 0)
    context_f.Enter()
    res_f = context_f.BuildCondBranch(fn2)
    context_t.ExitResult(res_f)
    context_f.Exit()

    # Add the final merge to the graph.
    if len(res_t) != len(res_f):
      raise ValueError("fn1 and fn2 must return the same number of tensors.")
    for x, y in zip(res_f, res_t):
      assert ((isinstance(x, ops.IndexedSlices) and
               isinstance(y, ops.IndexedSlices)) or
              (isinstance(x, ops.Tensor) and isinstance(y, ops.Tensor)))
      val_x = x if isinstance(x, ops.Tensor) else x.values
      val_y = y if isinstance(y, ops.Tensor) else y.values
      if val_x.dtype.base_dtype != val_y.dtype.base_dtype:
        raise ValueError("Outputs of fn1 and fn2 must have the same type: "
                         "%s, %s" % (val_x.dtype.name, val_y.dtype.name))
    merges = [merge([x[0], x[1]])[0] for x in zip(res_f, res_t)]
    return merges[0] if len(merges) == 1 else merges


# TODO(yuanbyu): We should probably separate the notion of context so it
# could be used not only for conditionals and loops but also subgraphs.
class WhileContext(ControlFlowContext):
  """The context for the loop construct."""

  def __init__(self, parallel_iterations, back_prop, name):
    self._name = ops.get_default_graph().unique_name(name)
    self._parallel_iterations = parallel_iterations
    self._back_prop = back_prop
    self._outer_context = None
    # We use this node to control constants created by the pred lambda.
    self._pivot_for_pred = None
    # We use this node to control constants created by the body lambda.
    self._pivot_for_body = None
    # The boolean tensor for loop termination condition. Used in code
    # generation for gradient computation
    self._pivot = None

    # The tensors for the counters added by AddForwardCounterLoop or
    # AddBackPropCounterLoop
    self._index = None

    # Information needed by backprop
    self._grad_context = None
    self._total_iterations = None
    self._history_map = {}
    self._switch_map = {}

    # values considered to have been already seen in this context
    self._values = set()

    # values referenced by but external to this context
    self._external_values = {}

  @property
  def name(self):
    return self._name

  @property
  def parallel_iterations(self):
    """The number of iterations allowed to run in parallel."""
    return self._parallel_iterations

  @property
  def back_prop(self):
    """True iff backprop is enabled for this While loop."""
    return self._back_prop

  @property
  def pivot(self):
    """The boolean tensor representing the loop termination condition."""
    return self._pivot

  @property
  def index(self):
    """The loop index representing the current iteration."""
    return self._index

  @property
  def grad_context(self):
    """The corresponding WhileContext for gradient."""
    return self._grad_context

  @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 in the While loop."""
    return self._switch_map

  @property
  def total_iterations(self):
    """The total number of iterations of the while loop."""
    return self._total_iterations

  def GetWhileContext(self):
    return self

  def GetControlPivot(self):
    if self._pivot_for_body:
      return self._pivot_for_body
    return self._pivot_for_pred

  def AddValue(self, val):
    """Add 'val' to the current context and its outer context recursively."""
    result = val
    if val.name not in self._values:
      self._values.add(val.name)
      if self._outer_context is not None:
        result = self._outer_context.AddValue(val)
      # Create an Enter that makes 'result' known to this context.
      enter = _Enter(result, self._name, is_constant=True,
                     parallel_iterations=self._parallel_iterations)
      self._values.add(enter.name)
      self._external_values[val.name] = enter
      result = enter
    else:
      actual_val = self._external_values.get(val.name)
      if actual_val is not None:
        result = actual_val
    return result

  def AddOp(self, op):
    """Adds 'op' to the current context."""
    if not op.inputs:
      if not op.control_inputs:
        # Add a control edge from the control pivot to this op.
        # pylint: disable=protected-access
        op._add_control_input(self.GetControlPivot().op)
        # pylint: enable=protected-access
      else:
        # Control edges must be in the same context.
        for x in op.control_inputs:
          assert x._get_control_flow_context() == self, (
              "Control inputs must come from Operations in the same while "
              "loop context (not an outer context).")
      for x in op.outputs:
        self._values.add(x.name)
    else:
      for index in range(len(op.inputs)):
        x = op.inputs[index]
        self.AddValue(x)
        real_x = self._external_values.get(x.name)
        if real_x is not None:
          op._update_input(index, real_x)
          # Add a control dependency to prevent loop invariants from
          # enabling ops that should not be executed.
          if real_x.op.type == "RefEnter" and real_x.op.get_attr("is_constant"):
            # pylint: disable=protected-access
            op._add_control_input(self.GetControlPivot().op)
            # pylint: enable=protected-access
      for x in op.outputs:
        self._values.add(x.name)

  def CreateGradWhileContext(self):
    """Creates the WhileContext for backprop gradient computation."""
    if self._grad_context is None:
      cnt = self.AddForwardCounterLoop()
      self._grad_context = WhileContext(self._parallel_iterations,
                                        self._back_prop, self._name)
      self._grad_context.AddBackPropCounterLoop(cnt)
    return self._grad_context

  def AddForwardCounterLoop(self):
    """Adds a loop that counts the number of iterations.

    This is added to the forward loop at the time when we start to
    create the loop for backprop gradient computation.

    The pseudocode is:
      `n = 0; while (_pivot) { n++; }`

    Returns:
      The number of iterations taken by the forward loop.
    """
    n = constant_op.constant(0, name="f_count")
    self.Enter()
    self.AddName(n.name)
    enter_n = _Enter(n, self._name, is_constant=False,
                     parallel_iterations=self._parallel_iterations,
                     name="f_count")
    merge_n = merge([enter_n, enter_n])[0]
    switch_n = switch(merge_n, self._pivot)
    self._index = switch_n[1]

    add_n = math_ops.add(self._index, 1)
    next_n = next_iteration(add_n)
    merge_n.op._update_input(1, next_n)

    self._total_iterations = exit(switch_n[0], name="f_count")
    self.Exit()
    return self._total_iterations

  def AddForwardAccumulateLoop(self, value):
    """Add an accumulation loop for each value needed in backprop.

    This is added to the forward loop at the first time when a value
    in the forward loop is used by backprop gradient computation loop.

    The pseudocode is:
    ```
      acc;
      while (_pivot) {
        if (index == 0) [value] else Concat(acc, [value]);
      }
    ```

    Args:
      value: The tensor that is accumulated.

    Returns:
      The accumulated history of value.

    Raises:
      ValueError: If the shape of "value" is not known statically.
    """
    if not value.get_shape().is_fully_defined():
      raise ValueError("Must have known shape: %s" % value)
    self._grad_context.Exit()
    # TODO(irving): Now that acc starts out empty, most of the
    # conditional logic can go away.
    acc = constant_op.constant([],
                               value.dtype,
                               shape=[0] + value.get_shape().as_list(),
                               name="f_acc")
    self.Enter()
    self.AddName(acc.name)
    enter_acc = _Enter(acc, self._name, is_constant=False,
                       parallel_iterations=self._parallel_iterations,
                       name="f_acc")
    merge_acc = merge([enter_acc, enter_acc])[0]
    switch_acc = switch(merge_acc, self._pivot)

    # If index = 0 then [value] else Concat(acc, [value]).
    cond = math_ops.greater(self._index, 0)
    switch_add_acc = switch(switch_acc[1], cond)
    expand_value = array_ops.expand_dims(value, 0)
    true_branch = array_ops.concat(0, [switch_add_acc[1], expand_value])
    false_branch = array_ops.identity(switch_add_acc[0])
    false_branch = with_dependencies([false_branch], expand_value)
    add_acc = merge([false_branch, true_branch])[0]

    next_acc = next_iteration(add_acc)
    merge_acc.op._update_input(1, next_acc)

    exit_acc = exit(switch_acc[0], name="f_acc")
    self.Exit()
    self._grad_context.Enter()
    return exit_acc

  def AddForwardAccumulateCondLoop(self, value):
    """Add an accumulation loop for each conditional switch.

    This is added to the forward loop at the first time when a conditional
    switch in the forward loop is used by backprop gradient computation loop.

    The pseudocode is:
      ```
      acc;
      while (_pivot) {
        Concat(acc, value);
      }
      ```

    Args:
      value: The boolean tensor that is accumulated.

    Returns:
      The accumulated history of value.
    """
    self._grad_context.Exit()
    acc = constant_op.constant(False, name="f_acc")
    self.Enter()
    self.AddName(acc.name)
    enter_acc = _Enter(acc, self._name, is_constant=False,
                       parallel_iterations=self._parallel_iterations,
                       name="f_acc")
    merge_acc = merge([enter_acc, enter_acc])[0]
    switch_acc = switch(merge_acc, self._pivot)
    acc = array_ops.concat(0, [switch_add_acc[1], value])
    next_acc = next_iteration(acc)
    merge_acc.op._update_input(1, next_acc)

    exit_acc = exit(switch_acc[0], name="f_acc")
    self.Exit()
    self._grad_context.Enter()
    return exit_acc

  def AddBackPropCounterLoop(self, count):
    """Add the backprop loop that controls the iterations.

    This is added to the backprop loop. It is used to control the loop
    termination and the slice index.

    The pseudocode is:
      `n = count; while (n >= 1) { n--; }`

    Args:
      count: The number of iterations for backprop.

    Returns:
      always 0.
    """
    one = constant_op.constant(1, name="b_count")
    self.Enter()
    self.AddName(count.name)
    enter_count = _Enter(count, self._name, is_constant=False,
                         parallel_iterations=self._parallel_iterations,
                         name="b_count")
    merge_count = merge([enter_count, enter_count])[0]
    self._pivot_for_pred = merge_count

    cond = math_ops.greater_equal(merge_count, one)
    self._pivot = loop_cond(cond, name="b_count")
    switch_count = switch(merge_count, self._pivot)

    # Add next_iteration right after Switch to match the gradient function.
    next_count = next_iteration(switch_count[1])
    self._pivot_for_body = next_count
    self._index = math_ops.sub(next_count, one)
    merge_count.op._update_input(1, self._index)

    exit_count = exit(switch_count[0], name="b_count")
    self.Exit()
    return exit_count

  def AddBackPropAccumulateLoop(self, value):
    """Add an accumulation loop for every loop invariant.

    This is added to the backprop loop. It is used to accumulate partial
    gradients for each loop iteration. Called when in the while context
    for gradient.

    The pseudocode is:
      ```
      acc = 0;
      while (_pivot) {
        acc += value;
      }
      ```

    Args:
      value: The partial gradient of an iteration for a loop invariant.

    Returns:
      The gradient for a loop invariant.
    """
    self.Exit()
    acc = constant_op.constant(0, value.dtype, name="b_acc")
    self.Enter()
    self.AddName(acc.name)
    enter_acc = _Enter(acc, self._name, is_constant=False,
                       parallel_iterations=self._parallel_iterations,
                       name="b_acc")
    merge_acc = merge([enter_acc, enter_acc], name="b_acc")[0]
    switch_acc = switch(merge_acc, self._pivot)

    next_acc = next_iteration(switch_acc[1])
    add_acc = math_ops.add(next_acc, value)
    merge_acc.op._update_input(1, add_acc)

    exit_acc = exit(switch_acc[0], name="b_acc")
    return exit_acc

  def BuildLoop(self, pred, body, loop_vars):
    """Add the loop termination condition and body to the graph."""

    loop_vars = ops.convert_n_to_tensor_or_indexed_slices(loop_vars)
    # Let the context know the loop variabes so the _Enter nodes below
    # would be added into the context correctly.
    self._values = set([x.name for x in loop_vars])
    if self._outer_context is not None:
      real_vars = [self._outer_context.AddValue(x) for x in loop_vars]
    else:
      real_vars = loop_vars
    enter_vars = [_Enter(x, self._name, is_constant=False,
                         parallel_iterations=self._parallel_iterations)
                  for x in real_vars]
    self._values = set([x.name for x in enter_vars])

    merge_vars = [merge([x, x])[0] for x in enter_vars]
    self._pivot_for_pred = merge_vars[0]

    # Build the graph for pred.
    c = ops.convert_to_tensor(pred(*merge_vars))
    self._pivot = loop_cond(c, name="LoopCond")
    switch_vars = [_SwitchRefOrTensor(x, self._pivot) for x in merge_vars]

    # Build the graph for body.
    vars_for_body = [_Identity(x[1]) for x in switch_vars]
    self._pivot_for_body = vars_for_body[0]

    body_result = body(*vars_for_body)
    if not isinstance(body_result, (list, _basetuple)):
      body_result = [body_result]
    result = ops.convert_n_to_tensor_or_indexed_slices(body_result)
    next_vars = [next_iteration(x) for x in result]

    # Add the back edges to complete the loop.
    assert len(merge_vars) == len(next_vars)
    for x in zip(merge_vars, next_vars):
      x[0].op._update_input(1, x[1])

    # Add the exit ops.
    exit_vars = [exit(x[0]) for x in switch_vars]

    for m_var, n_var, e_var in zip(merge_vars, next_vars, exit_vars):
      if m_var.get_shape().is_compatible_with(n_var.get_shape()):
        e_var.set_shape(m_var.get_shape().merge_with(n_var.get_shape()))

    # Exit the loop.
    self.ExitResult(exit_vars)
    self.Exit()
    return exit_vars[0] if len(exit_vars) == 1 else exit_vars


def While(cond, body, loop_vars, parallel_iterations=10, back_prop=True,
          name=None):
  """Repeat `body` while the condition `cond` is true.

  `cond` is a function taking a list of tensors and returning a boolean scalar
  tensor. `body` is a function taking a list of tensors and returning a list of
  tensors of the same length and with the same types as the input. `loop_vars`
  is a list of tensors that is passed to both `cond` and `body`.

  While `cond` evaluates to true, `body` is executed.

  Args:
    cond: The termination condition of the loop.
    body: A function that represents the loop body.
    loop_vars: The list of variable input tensors.
    parallel_iterations: The number of iterations allowed to run in parallel.
    back_prop: Whether backprop is enabled for this while loop.
    name: Optional name prefix for the returned tensors.

  Returns:
    The output tensors for the loop variables after the loop.

  Raises:
    TypeError: if `cond` or `body` is not callable.
    ValueError: if `loop_var` is empty.

  Example:
    ```python
    i =  Constant(0)
    c = lambda i: math_ops.less(i, 10)
    b = lambda i: math_ops.add(i, 1)
    r = While(c, b, [i])
    ```
  """
  with ops.op_scope(loop_vars, name, "While") as name:
    if not loop_vars:
      raise ValueError("No loop variables provided")
    if not callable(cond):
      raise TypeError("cond must be callable.")
    if not callable(body):
      raise TypeError("body must be callable.")

    context = WhileContext(parallel_iterations, back_prop, name)
    context.Enter()
    return context.BuildLoop(cond, body, loop_vars)


def _AsTensorList(x, p):
  """Return x as a list of Tensors or IndexedSlices.

  For entries of `x` that are Operations, this returns an Identity of `p`
  with a dependency on the operation.

  Args:
    x: A Tensor/IndexedSlices/Operation or a list or tuple of them.
    p: A Tensor to return for entries in `x` that are Operations.

  Returns:
    A list of Tensors or IndexedSlices.
  """
  if not isinstance(x, list) and not isinstance(x, _basetuple):
    x = [x]

  l = []
  for v in x:
    if isinstance(v, ops.Operation):
      v = with_dependencies([v], p)
    v = ops.convert_to_tensor_or_indexed_slices(v)
    if isinstance(v, ops.Tensor):
      l.append(array_ops.identity(v))
    else:
      l.append(ops.IndexedSlices(array_ops.identity(v.values),
                                 array_ops.identity(v.indices)))
  return l


def _CheckResults(a, b):
  assert len(a) == len(b), (
      "Values returned by a() and b() must have the same length.")
  for x, y in zip(a, b):
    assert x.dtype == y.dtype, (
        "Values returned by a() [%s] and b() [%s] must have "
        "the same type: %s, %s." %
        (x.name, y.name, x.dtype.name, y.dtype.name))


def with_dependencies(dependencies, output_tensor, name=None):
  """Produces the content of `output_tensor` only after `dependencies`.

  In some cases, a user may want the output of an operation to be
  consumed externally only after some other dependencies have run
  first. This function ensures returns `output_tensor`, but only after all
  operations in `dependencies` have run. Note that this means that there is
  no guarantee that `output_tensor` will be evaluated after any `dependencies`
  have run.

  See also `tuple` and `group`.

  Args:
    dependencies: A list of operations to run before this op finishes.
    output_tensor: A `Tensor` or `IndexedSlices` that will be returned.
    name: (Optional) A name for this operation.

  Returns:
    Same as `output_tensor`.

  Raises:
    TypeError: if `output_tensor` is not a `Tensor` or `IndexedSlices`.
  """
  with ops.op_scope(dependencies + [output_tensor], name,
                   "control_dependency") as name:
    with ops.device(output_tensor.device
                    or ops.get_default_graph().get_default_device()):
      with ops.control_dependencies(dependencies):
        output_tensor = ops.convert_to_tensor_or_indexed_slices(output_tensor)
        if isinstance(output_tensor, ops.Tensor):
          return _Identity(output_tensor, name=name)
        else:
          return ops.IndexedSlices(_Identity(output_tensor.values, name=name),
                                   output_tensor.indices,
                                   output_tensor.dense_shape)


def _GroupControlDeps(dev, deps, name=None):
  with ops.control_dependencies(deps):
    if dev is None:
      return no_op(name=name)
    else:
      with ops.device(dev):
        return no_op(name=name)


# TODO(mdevin): Accept "inputs" as a list.
def group(*inputs, **kwargs):
  """Create an op that groups multiple operations.

  When this op finishes, all ops in `input` have finished. This op has no
  output.

  See also `tuple` and `with_dependencies`.

  Args:
    *inputs: One or more tensors to group.
    **kwargs: Optional parameters to pass when constructing the NodeDef.
    name: A name for this operation (optional).

  Returns:
    An Operation that executes all its inputs.

  Raises:
    ValueError: If an unknown keyword argument is provided, or if there are
                no inputs.
  """
  name = kwargs.pop("name", None)
  if kwargs:
    raise ValueError("Unknown keyword arguments: " + ", ".join(kwargs.keys()))
  if not inputs:
    # TODO(mdevin): Would make sense to return a NoOp.
    raise ValueError("No inputs provided")
  with ops.op_scope(inputs, name, "group_deps") as name:
    # Sorts *inputs according to their devices.
    ops_on_device = {}  # device -> operations specified on the device.
    for inp in inputs:
      dev = inp.device
      if dev in ops_on_device:
        ops_on_device[dev].append(inp)
      else:
        ops_on_device[dev] = [inp]
    if len(ops_on_device) == 1:
      # 1-level tree. The root node is the returned NoOp node.
      dev, deps = ops_on_device.items()[0]
      return _GroupControlDeps(dev, deps, name=name)
    # 2-level tree. The root node is the returned NoOp node.
    # deps contains 1 NoOp node for each device.
    deps = []
    for dev in sorted(ops_on_device.iterkeys()):
      deps.append(_GroupControlDeps(dev, ops_on_device[dev]))
    return _GroupControlDeps(None, deps, name=name)

def tuple(tensors, name=None, control_inputs=None):
  """Group tensors together.

  This creates a tuple of tensors with the same values as the `tensors`
  argument, except that the value of each tensor is only returned after the
  values of all tensors have been computed.

  `control_inputs` contains additional ops that have to finish before this op
  finishes, but whose outputs are not returned.

  This can be used as a "join" mechanism for parallel computations: all the
  argument tensors can be computed in parallel, but the values of any tensor
  returned by `tuple` are only available after all the parallel computations
  are done.

  See also `group` and `with_dependencies`.

  Args:
    tensors: A list of `Tensor`s or `IndexedSlices`, some entries can be `None`.
    name: (optional) A name to use as a `name_scope` for the operation.
    control_inputs: List of additional ops to finish before returning.

  Returns:
    Same as `tensors`.

  Raises:
    ValueError: If `tensors` does not contain any `Tensor` or `IndexedSlices`.

  """
  with ops.op_scope(tensors, name, "tuple") as name:
    gating_ops = [t.op for t in tensors if t]
    if control_inputs:
      gating_ops += control_inputs
    # Note that in order to ensure ordering in the pbtxt, we must take care to
    # ensure the order here.
    gating_ops = sorted(set(gating_ops), key=lambda op: op._id)  # Uniquify ops.
    if not gating_ops:
      raise ValueError("Must have at least one Tensor: %s" % tensors)
    gate = group(*gating_ops)
    tpl = []
    for t in tensors:
      if t:
        tpl.append(with_dependencies([gate], t))
      else:
        tpl.append(None)
    return tpl


# TODO(yuanbyu): It would be nicer if we could have the distributed list
# support that Derek has been proposing.
# TODO(yuanbyu, mrry): Handle stride to support sliding windows.
def fold(fn, elems, elem_shape, name=None):
  """The fold operator on slices of a tensor.

  This fold operator applies the function `fn` to slices of `elems` on
  dimension 0. The shape of the slices is specified by `elem_shape`. `elems`
  must contain at least one slice (`shape(elems)[0] / elem_shape[0] > 0`).

  Args:
    fn: The function to be performed on each slice of the tensor.
    elems: The tensor to whose slices we want to apply `fn`.
    elem_shape: The shape definition for the slices.
    name: Optional name prefix for the returned tensors.

  Returns:
    A tensor resulting from applying `fn` consecutively on each slice of
    `elems`.

  Raises:
    TypeError: if `fn` is not callable.
  """
  with ops.op_scope([elems], name, "Fold") as name:
    if not callable(fn):
      raise TypeError("fn must be callable.")

    s0 = array_ops.shape(elems)[0]
    d0 = elem_shape[0]
    n = math_ops.div(s0, d0)
    b1 = array_ops.zeros(array_ops.expand_dims(array_ops.rank(elems) - 1, 0),
                         dtype=types.int32)
    # Initialize the output with slice 0
    b = array_ops.concat(0, [[0], b1])
    o = array_ops.slice(elems, b, elem_shape)
    i = ops.convert_to_tensor(d0)

    def Compute(i, o):
      b = array_ops.concat(0, [array_ops.expand_dims(i, 0), b1])
      x = array_ops.slice(elems, b, elem_shape)
      o = fn(o, x)
      i = math_ops.add(i, d0)
      return [i, o]
    r = While(lambda i, o: math_ops.less(i, n), Compute, [i, o])
    return r[1]


def case(pred_fn_pairs, default, exclusive=False, name="Case"):
  """Create a Case operation.

  The `pred_fn_pairs` parameter is a dict or list of pairs of size N.
  Each pair contains a boolean scalar tensor and a python callable that
  creates the tensors to be returned if the boolean evaluates to True. `default`
  is a callable generating a list of tensors. All the callables in
  `pred_fn_pairs` as well as `default` should return the same number and types
  of tensors.

  If `exclusive==True`, all predicates are evaluated, and a logging operation
  with an error is returned if more than one of the predicates evaluates to
  True. If `exclusive==False`, execution stops are the first predicate which
  evaluates to True, and the tensors generated by the corresponding function
  are returned immediately. If none of the predicates evaluate to True, this
  operation returns the tensors generated by `default`.

  Example 1:
    Pseudocode:
    ```
      if (x < y) return 17;
      else return 23;
    ```

    Expressions:
    ```
      f1 = lambda: Constant(17)
      f2 = lambda: Constant(23)
      r = Case([(math_ops.less(x, y), f1)], default=f2)
    ```

  Example 2:
    Pseudocode:
    ```
      if (x < y && x > z) raise OpError("Only one predicate may evaluate true");
      if (x < y) return 17;
      else if (x > z) return 23;
      else return -1;
    ```

    Expressions:
    ```
      def f1(): return Constant(17)
      def f2(): return Constant(23)
      def f3(): return Constant(-1)
      r = Case({math_ops.less(x, y): f1, math_ops.greater(x, z): f2},
               default=f3, exclusive=True)
    ```

  Args:
    pred_fn_pairs: Dict or list of pairs of a boolean scalar tensor and a
                   callable which returns a list of tensors.
    default: A callable that returns a list of tensors.
    exclusive: True iff more than one predicate is allowed to evaluate to True.
    name: A name for this operation (optional).

  Returns:
    The tensors returned by the first pair whose predicate evaluated to True, or
    those returned by `default` if none does.

  Raises:
    TypeError: If `pred_fn_pairs` is not a list/dictionary.
    TypeError: If `pred_fn_pairs` is a list but does not contain 2-tuples.
    TypeError: If `fns[i]` is not callable for any i, or `default` is not
               callable.
  """
  pfp = pred_fn_pairs  # For readability
  if not (isinstance(pfp, list) or isinstance(pfp, _basetuple)
          or isinstance(pfp, dict)):
    raise TypeError("fns must be a list, tuple, or dict")
  if isinstance(pfp, dict):
    pfp = pfp.items()
    if not exclusive:
      logging.warn("%s: Provided dictionary of predicate/fn pairs, but "
                   "exclusive=False.  Order of conditional tests is "
                   "not guaranteed." % name)
  for tup in pfp:
    if not isinstance(tup, _basetuple) or len(tup) != 2:
      raise TypeError("Each entry in pred_fn_pairs must be a 2-tuple")
    pred, fn = tup
    if pred.dtype != types.bool:
      raise TypeError("pred must be of type bool: %s", pred.name)
    if not callable(fn):
      raise TypeError("fn for pred %s must be callable." % pred.name)
  if not callable(default):
    raise TypeError("default must be callable.")

  preds, fns = map(list, zip(*pfp))
  with ops.op_scope([[f() for f in fns] + preds + [default()]], name, "Case"):
    if not preds:
      return default()
    not_preds = []
    for i, p in enumerate(preds):
      with ops.name_scope("not_%d" % i):
        not_preds.append(math_ops.logical_not(p))
    and_not_preds = [constant_op.constant(True, name="and_not_true")]
    for i, notp in enumerate(not_preds[:-1]):
      with ops.name_scope("and_not_%d" % i):
        and_not_preds.append(math_ops.logical_and(and_not_preds[-1], notp))

    # preds = [p1, p2, p3]
    # fns = [f1, f2, f3]
    # not_preds = [~p1, ~p2, ~p3]
    # case_preds = [p1 & True,
    #               p2 & ~p1,
    #               p3 & ~p1 & ~ p2]
    case_preds = []
    for i, (p, and_not_p_prev) in enumerate(zip(preds, and_not_preds)):
      with ops.name_scope("case_%d" % i):
        case_preds.append(math_ops.logical_and(p, and_not_p_prev))

    # case_sequence = [Cond(p3 & ..., f3, default),
    #                  Cond(p2 & ..., f2, lambda: case_sequence[0]),
    #                  ...
    #                  Cond(p1 & True, f1, lambda: case_sequence[i-1])]
    # and prev_case_seq will loop from case_sequence[0] to case_sequence[-1]
    if exclusive:
      # TODO(ebrevdo): Add Where() for DT_BOOL, replace with Size(Where(preds))
      preds_c = array_ops.concat(0, preds, name="preds_c")
      num_true_conditions = math_ops.reduce_sum(
          math_ops.cast(preds_c, types.int32), name="num_true_conds")
      at_most_one_true_condition = math_ops.less(
          num_true_conditions, constant_op.constant(2, name="two_true_conds"))

      error_msg = [
          ("More than one condition evaluated as True but "
           "exclusive=True.  Conditions: (%s), Values:"
           % ", ".join([p.name for p in preds])),
          preds_c]
      with ops.control_dependencies([
          logging_ops.Assert(condition=at_most_one_true_condition,
                             data=error_msg, summarize=len(preds))]):
        prev_case_seq = default()
        for i, (cp, fn) in enumerate(zip(case_preds, fns)[::-1]):
          prev_case_seq = cond(cp, fn, lambda: prev_case_seq, name="If_%d" % i)
    else:
      prev_case_seq = default()
      for i, (cp, fn) in enumerate(zip(case_preds, fns)[::-1]):
        prev_case_seq = cond(cp, fn, lambda: prev_case_seq, name="If_%d" % i)

    return prev_case_seq


ops.RegisterShape("Enter")(common_shapes.unchanged_shape)
ops.RegisterShape("Exit")(common_shapes.unknown_shape)
ops.RegisterShape("NextIteration")(common_shapes.unchanged_shape)
ops.RegisterShape("RefEnter")(common_shapes.unchanged_shape)
ops.RegisterShape("ControlTrigger")(common_shapes.no_outputs)
ops.RegisterShape("NoOp")(common_shapes.no_outputs)


@ops.RegisterShape("LoopCond")
def _LoopCondShape(op):
  """Shape function for the LoopCond op."""
  return [op.inputs[0].get_shape().merge_with(tensor_shape.scalar())]


@ops.RegisterShape("Merge")
def _MergeShape(op):
  """Shape function for the Merge op.

  The Merge op takes many inputs of arbitrary shapes, and produces a
  first output that is one of those inputs, and a second scalar
  output.

  This function conservatively assumes that if any of its inputs is
  not fully defined, the output shape is unknown. If all of the inputs
  have the exact same known shape, the output must have that shape.

  Args:
    op: A Merge Operation.

  Returns:
    A single-element list containing the Shape of the Merge op.

  """
  first_input_shape = op.inputs[0].get_shape()
  if first_input_shape.is_fully_defined():
    for input_ in op.inputs[1:]:
      input_shape = input_.get_shape()
      if (not input_shape.is_fully_defined()
          or not input_shape.is_compatible_with(first_input_shape)):
        return [tensor_shape.unknown_shape(), tensor_shape.scalar()]
    return [first_input_shape, tensor_shape.scalar()]
  else:
    return [tensor_shape.unknown_shape(), tensor_shape.scalar()]


@ops.RegisterShape("RefSelect")
def _RefSelectShape(op):
  """Shape function for the RefSelect op.

  The RefSelect takes one scalar input and N inputs of arbitrary
  shapes, and produces one output, which is one of those N inputs.

  This function conservatively assumes that if any of the N inputs is
  not fully defined, the output shape is unknown. If all of the N
  inputs have the exact same known shape, the output must have that
  shape.

  Args:
    op: A RefSelect Operation.

  Returns:
    A single-element list containing the Shape of the RefSelect op.
  """
  unused_shape = op.inputs[0].get_shape().merge_with(tensor_shape.scalar())
  first_input_shape = op.inputs[1].get_shape()
  if first_input_shape.is_fully_defined():
    for input_ in op.inputs[2:]:
      input_shape = input_.get_shape()
      if (not input_shape.is_fully_defined()
          or not input_shape.is_compatible_with(first_input_shape)):
        return [tensor_shape.unknown_shape()]
    return [first_input_shape]
  else:
    return [tensor_shape.unknown_shape()]


@ops.RegisterShape("RefSwitch")
@ops.RegisterShape("Switch")
def _SwitchShape(op):
  input_shape = op.inputs[0].get_shape()
  unused_pred_shape = op.inputs[1].get_shape().merge_with(tensor_shape.scalar())
  return [input_shape] * 2