util_test.py 文件源码

def test_make_list_of_t(self):
    """Test for ge.util.make_list_of_t."""
    g0 = ops.Graph()
    with g0.as_default():
      a0 = constant_op.constant(1)
      b0 = constant_op.constant(2)
      c0 = math_ops.add(a0, b0)  # pylint: disable=unused-variable
    # Should extract the tensors from tre graph.
    self.assertEqual(len(ge.util.make_list_of_t(g0)), 3)
    # Should extract the tensors from the tuple
    self.assertEqual(len(ge.util.make_list_of_t((a0, b0))), 2)
    # Should extract the tensors and ignore the ops.
    self.assertEqual(
        len(ge.util.make_list_of_t(
            (a0, a0.op, b0), ignore_ops=True)), 2)