def testStack(self):
   l = list_ops.empty_tensor_list(element_dtype=dtypes.float32,
                                  element_shape=scalar_shape())
   l = list_ops.tensor_list_push_back(l, constant_op.constant(1.0))
   l = list_ops.tensor_list_push_back(l, constant_op.constant(2.0))
   t = list_ops.tensor_list_stack(l, element_dtype=dtypes.float32)
   self.assertAllEqual(t, [1.0, 2.0])

 def testUnknownShape(self):
   l = list_ops.empty_tensor_list(element_dtype=dtypes.float32,
                                  element_shape=-1)
   l = list_ops.tensor_list_push_back(l, constant_op.constant(1.0))
   l = list_ops.tensor_list_push_back(l, constant_op.constant([1.0, 2.0]))
   _, e = list_ops.tensor_list_pop_back(l, element_dtype=dtypes.float32)
   self.assertAllEqual(e, [1.0, 2.0])

 def testConcatWithNonFullyDefinedElementShape(self):
   l = list_ops.empty_tensor_list(
       element_dtype=dtypes.float32, element_shape=[None, 2])
   l = list_ops.tensor_list_push_back(l, [[0., 1.]])
   l = list_ops.tensor_list_push_back(l, [[2., 3.], [4., 5.]])
   t = list_ops.tensor_list_concat(l, element_dtype=dtypes.float32)
   self.assertAllEqual(self.evaluate(t), [[0., 1.], [2., 3.], [4., 5.]])

 def testPushInFullListFails(self):
   l = list_ops.empty_tensor_list(
       element_dtype=dtypes.float32, element_shape=[], max_num_elements=1)
   l = list_ops.tensor_list_push_back(l, constant_op.constant(1.0))
   with self.assertRaisesRegexp(errors.InvalidArgumentError,
                                "Tried to push item into a full list"):
     l = list_ops.tensor_list_push_back(l, 2.)
     self.evaluate(l)

 def _testStack(self, max_num_elements):
   l = list_ops.empty_tensor_list(
       element_dtype=dtypes.float32,
       element_shape=scalar_shape(),
       max_num_elements=max_num_elements)
   l = list_ops.tensor_list_push_back(l, constant_op.constant(1.0))
   l = list_ops.tensor_list_push_back(l, constant_op.constant(2.0))
   t = list_ops.tensor_list_stack(l, element_dtype=dtypes.float32)
   self.assertAllEqual(self.evaluate(t), [1.0, 2.0])

 def testPushInEmptyListWithUnknownElementShape(self):
   with self.cached_session(), self.test_scope():
     l = list_ops.empty_tensor_list(
         element_dtype=dtypes.float32, element_shape=None, max_num_elements=2)
     l = list_ops.tensor_list_push_back(l, [3.0, 4.0])
     # Pushing an element with a different shape should raise an error.
     with self.assertRaisesRegexp(errors.InvalidArgumentError, "Shape"):
       l = list_ops.tensor_list_push_back(l, 5.)
       self.evaluate(
           list_ops.tensor_list_stack(l, element_dtype=dtypes.float32))

 def testSetDoesNotUpdatePushIndex(self):
   with self.cached_session(), self.test_scope():
     l = list_ops.empty_tensor_list(
         element_shape=[], element_dtype=dtypes.float32, max_num_elements=2)
     # SetItem should not change the push index.
     l = list_ops.tensor_list_set_item(l, 1, 3.)
     l = list_ops.tensor_list_push_back(l, 5.)
     l = list_ops.tensor_list_push_back(l, 7.)
     t = list_ops.tensor_list_stack(l, element_dtype=dtypes.float32)
     self.assertAllEqual(t, [5., 7.])

 def _testStack(self, max_num_elements):
   l = list_ops.empty_tensor_list(
       element_dtype=dtypes.float32,
       element_shape=[],
       max_num_elements=max_num_elements)
   l = list_ops.tensor_list_push_back(l, constant_op.constant(1.0))
   l = list_ops.tensor_list_push_back(l, constant_op.constant(2.0))
   t = list_ops.tensor_list_stack(l, element_dtype=dtypes.float32)
   if not context.executing_eagerly():
     self.assertAllEqual(t.shape.as_list(), [None])
   self.assertAllEqual(self.evaluate(t), [1.0, 2.0])

 def testConcatWithMismatchingTensorShapesFails(self):
   l = list_ops.empty_tensor_list(
       element_dtype=dtypes.float32, element_shape=None)
   l = list_ops.tensor_list_push_back(l, [[0., 1.]])
   l = list_ops.tensor_list_push_back(l, [[2.], [4.]])
   with self.assertRaisesRegexp(
       errors.InvalidArgumentError,
       r"Tried to concat tensors with unequal shapes: "
       r"\[2\] vs \[1\]"):
     t = list_ops.tensor_list_concat(l, element_dtype=dtypes.float32)
     self.evaluate(t)

 def testStack(self):
   with self.cached_session(), self.test_scope():
     l = list_ops.empty_tensor_list(
         element_dtype=dtypes.float32,
         element_shape=[],
         max_num_elements=2)
     l = list_ops.tensor_list_push_back(l, constant_op.constant(1.0))
     e = list_ops.tensor_list_get_item(l, 0, element_dtype=dtypes.float32)
     self.assertAllEqual(e, 1.0)
     l = list_ops.tensor_list_push_back(l, constant_op.constant(2.0))
     t = list_ops.tensor_list_stack(l, element_dtype=dtypes.float32)
     self.assertAllEqual(t.shape.as_list(), [None])
     self.assertAllEqual(t, [1.0, 2.0])

 def testPushPop(self):
   with self.cached_session() as sess, self.test_scope():
     num = array_ops.placeholder(dtypes.int32)
     l = list_ops.tensor_list_reserve(
         element_shape=(7, 15), num_elements=num, element_dtype=dtypes.float32)
     l = list_ops.tensor_list_push_back(
         l, constant_op.constant(1.0, shape=(7, 15)))
     l = list_ops.tensor_list_push_back(
         l, constant_op.constant(2.0, shape=(7, 15)))
     l, e2 = list_ops.tensor_list_pop_back(l, element_dtype=dtypes.float32)
     _, e1 = list_ops.tensor_list_pop_back(l, element_dtype=dtypes.float32)
     self.assertAllEqual(sess.run(e2, {num: 10}), 2.0 * np.ones((7, 15)))
     self.assertAllEqual(sess.run(e1, {num: 10}), 1.0 * np.ones((7, 15)))

 def testGatherGrad(self):
   with backprop.GradientTape() as tape:
     l = list_ops.empty_tensor_list(element_dtype=dtypes.float32,
                                    element_shape=scalar_shape())
     c0 = constant_op.constant(1.0)
     tape.watch(c0)
     l = list_ops.tensor_list_push_back(l, c0)
     l = list_ops.tensor_list_push_back(l, constant_op.constant(2.0))
     t = list_ops.tensor_list_gather(l, [1, 0], element_dtype=dtypes.float32)
     self.assertAllEqual(self.evaluate(t), [2.0, 1.0])
     s = (t[0] + t[1]) * (t[0] + t[1])
   dt = tape.gradient(s, c0)
   self.assertAllEqual(self.evaluate(dt), 6.0)

 def testPushPop(self):
   with self.cached_session() as sess, self.test_scope():
     l = list_ops.empty_tensor_list(
         element_shape=(7, 15),
         element_dtype=dtypes.float32,
         max_num_elements=10)
     l = list_ops.tensor_list_push_back(
         l, constant_op.constant(1.0, shape=(7, 15)))
     l = list_ops.tensor_list_push_back(
         l, constant_op.constant(2.0, shape=(7, 15)))
     l, e2 = list_ops.tensor_list_pop_back(l, element_dtype=dtypes.float32)
     _, e1 = list_ops.tensor_list_pop_back(l, element_dtype=dtypes.float32)
     self.assertAllEqual(sess.run(e2), 2.0 * np.ones((7, 15)))
     self.assertAllEqual(sess.run(e1), 1.0 * np.ones((7, 15)))

  def testStackWithPartiallyDefinedElementShape(self):
    l = list_ops.empty_tensor_list(
        element_dtype=dtypes.float32, element_shape=[-1])
    l = list_ops.tensor_list_push_back(l, constant_op.constant([1.0]))
    l = list_ops.tensor_list_push_back(l, constant_op.constant([2.0]))

    t = list_ops.tensor_list_stack(l, element_dtype=dtypes.float32)
    self.assertAllEqual(self.evaluate(t), [[1.0], [2.0]])

    # Should raise an error when the element tensors do not all have the same
    # shape.
    with self.assertRaisesRegexp(errors.InvalidArgumentError, "unequal shapes"):
      l = list_ops.tensor_list_push_back(l, constant_op.constant([2.0, 3.0]))
      t = list_ops.tensor_list_stack(l, element_dtype=dtypes.float32)
      self.evaluate(t)

def _WhileGrad(op, *grads):  # pylint: disable=invalid-name
  """The gradient of a While op produced by while_loop."""
  body_graph = _get_body_graph(op)

  # Replace None gradients with zeros. This is needed because `grads` could have
  # None incoming gradients for the TensorLists. If we pass None's through, the
  # custom gradient of TensorListPopBack will create an EmptyTensorList inside
  # the FuncGraph which is undesirable.
  # TODO(b/80444525): There might be an issue with treating no gradient as zero
  # gradient in certain cases. Consider replacing None gradients with Zeros
  # for accumulators only.
  grads = [
      g if g is not None else array_ops.zeros_like(output)
      for g, output in zip(grads, op.outputs)
  ]

  body_grad_graph, args = _create_grad_func(
      body_graph, grads,
      util.unique_grad_fn_name(body_graph.name), op)

  intermediate_tensors = _get_intermediates(body_grad_graph)

  for intermediate_tensor in intermediate_tensors:
    tensor_list = list_ops.empty_tensor_list(
        element_dtype=intermediate_tensor.dtype,
        element_shape=_get_tensor_convertible_shape(intermediate_tensor.shape))
    with body_grad_graph.as_default():
      tensor_list_ph = body_grad_graph.capture(tensor_list, whitelisted=True)
      # Push the intermediate tensor to the tensor list.
      appended_tensor_list = list_ops.tensor_list_push_back(tensor_list_ph,
                                                            intermediate_tensor)
      # Add this modified tensor list to the list of outputs.
      body_grad_graph.outputs.append(appended_tensor_list)

  def grad_cond(counter, max_iters, *unused_args):
    return counter < max_iters

  loop_vars = args + body_grad_graph.external_captures
  grad_cond_name = util.unique_grad_fn_name(op.get_attr("cond").name)
  cond_grad_graph = func_graph_module.func_graph_from_py_func(
      grad_cond_name, grad_cond, loop_vars, {},
      func_graph=util.WhileCondFuncGraph(grad_cond_name))

  assert len(loop_vars) == len(body_grad_graph.inputs)
  assert len(loop_vars) == len(body_grad_graph.outputs)
  assert len(loop_vars) == len(cond_grad_graph.inputs)

  outputs = gen_functional_ops._while(
      loop_vars,
      util.create_new_tf_function(cond_grad_graph),
      util.create_new_tf_function(body_grad_graph),
      output_shapes=[t.shape for t in body_grad_graph.outputs],
      name="%s_grad" % op.name)

  _copy_handle_data(body_grad_graph.outputs, outputs)
  _maybe_set_lowering_attr(outputs[0].op)

  # outputs[0] is the loop counter.
  # outputs[1] is the total number of loop iterations.
  return outputs[2:2 + len(op.inputs)]

 def testConcatListWithScalarElementsFails(self):
   l = list_ops.empty_tensor_list(
       element_dtype=dtypes.float32, element_shape=None)
   l1 = list_ops.tensor_list_push_back(l, 1.)
   with self.assertRaisesRegexp(
       errors.InvalidArgumentError, "Concat saw a scalar shape at index 0"
       " but requires at least vectors"):
     t = list_ops.tensor_list_concat(l1, element_dtype=dtypes.float32)
     self.evaluate(t)
   l1 = list_ops.tensor_list_push_back(l, [1.])
   l1 = list_ops.tensor_list_push_back(l1, 2.)
   with self.assertRaisesRegexp(
       errors.InvalidArgumentError, "Concat saw a scalar shape at index 1"
       " but requires at least vectors"):
     t = list_ops.tensor_list_concat(l1, element_dtype=dtypes.float32)
     self.evaluate(t)

  def testDefaultGradYs(self):
    with ops.Graph().as_default():
      tl = list_ops.empty_tensor_list(
          element_dtype=dtypes.float32,
          element_shape=ops.convert_to_tensor([], dtype=dtypes.int32))
      a = constant(1.0)
      tl = list_ops.tensor_list_push_back(tl, a)

      grad_tl = list_ops.empty_tensor_list(
          element_dtype=dtypes.float32,
          element_shape=ops.convert_to_tensor([], dtype=dtypes.int32))
      grad_tl = list_ops.tensor_list_push_back(tl, constant(5.0))

      grad = gradients.gradients(tl, a, grad_ys=grad_tl)[0]
      with self.cached_session() as sess:
        self.assertEquals(self.evaluate(grad), 5.)

 def testPushPopSeparateLists(self):
   with self.cached_session() as sess, self.test_scope():
     l = list_ops.empty_tensor_list(
         element_shape=[],
         element_dtype=dtypes.float32,
         max_num_elements=20)
     l = list_ops.tensor_list_push_back(l, constant_op.constant(1.0))
     l2 = list_ops.tensor_list_push_back(l, constant_op.constant(2.0))
     l3 = list_ops.tensor_list_push_back(l, constant_op.constant(3.0))
     _, e11 = list_ops.tensor_list_pop_back(l, element_dtype=dtypes.float32)
     l2, e21 = list_ops.tensor_list_pop_back(l2, element_dtype=dtypes.float32)
     l2, e22 = list_ops.tensor_list_pop_back(l2, element_dtype=dtypes.float32)
     l3, e31 = list_ops.tensor_list_pop_back(l3, element_dtype=dtypes.float32)
     l3, e32 = list_ops.tensor_list_pop_back(l3, element_dtype=dtypes.float32)
     result = sess.run([e11, [e21, e22], [e31, e32]])
     self.assertEqual(result, [1.0, [2.0, 1.0], [3.0, 1.0]])

 def testPushPopSeparateLists(self):
   with self.cached_session() as sess, self.test_scope():
     num = array_ops.placeholder(dtypes.int32)
     l = list_ops.tensor_list_reserve(
         element_shape=scalar_shape(),
         num_elements=num,
         element_dtype=dtypes.float32)
     l = list_ops.tensor_list_push_back(l, constant_op.constant(1.0))
     l2 = list_ops.tensor_list_push_back(l, constant_op.constant(2.0))
     l3 = list_ops.tensor_list_push_back(l, constant_op.constant(3.0))
     _, e11 = list_ops.tensor_list_pop_back(l, element_dtype=dtypes.float32)
     l2, e21 = list_ops.tensor_list_pop_back(l2, element_dtype=dtypes.float32)
     l2, e22 = list_ops.tensor_list_pop_back(l2, element_dtype=dtypes.float32)
     l3, e31 = list_ops.tensor_list_pop_back(l3, element_dtype=dtypes.float32)
     l3, e32 = list_ops.tensor_list_pop_back(l3, element_dtype=dtypes.float32)
     result = sess.run([e11, [e21, e22], [e31, e32]], {num: 20})
     self.assertEqual(result, [1.0, [2.0, 1.0], [3.0, 1.0]])

      def body(i, m, t1):
        t1 = control_flow_ops.cond(
            math_ops.equal(list_ops.tensor_list_length(t1), 0),
            lambda: list_ops.empty_tensor_list(m.shape, m.dtype), lambda: t1)

        t1 = list_ops.tensor_list_push_back(t1, m * i)
        i += 1.0
        return i, m, t1