| 925 | @test_util.enable_control_flow_v2 |
| 926 | @test_util.run_v1_only("b/120545219") |
| 927 | def testAddWhileInput(self): |
| 928 | if forward_compat.forward_compatible(2019, 8, 23): |
| 929 | @eager_function.defun |
| 930 | def test(): |
| 931 | output = control_flow_ops.while_loop(lambda x: x < 3, lambda x: x + 1, |
| 932 | [1]) |
| 933 | while_op = output.op |
| 934 | self.assertEqual(while_op.type, "StatelessWhile") |
| 935 | orig_num_inputs = len(while_op.inputs) |
| 936 | |
| 937 | # Make sure we can handle the while op having a control input. |
| 938 | while_op._add_control_input(constant_op.constant(0).op) |
| 939 | |
| 940 | new_input1 = constant_op.constant(1.0) |
| 941 | new_input2 = constant_op.constant(True) |
| 942 | |
| 943 | # Clear output shapes to bypass shape checking. |
| 944 | while_op._set_shape_list_attr("output_shapes", []) |
| 945 | while_op._set_type_list_attr("T", |
| 946 | [t.dtype for t in while_op.inputs] + |
| 947 | [new_input1.dtype, new_input2.dtype]) |
| 948 | |
| 949 | while_op._add_while_inputs([new_input1, new_input2]) |
| 950 | # Can't add an edge beyond what's specified by "T" |
| 951 | with self.assertRaises(errors.OutOfRangeError): |
| 952 | while_op._add_while_inputs([new_input2]) |
| 953 | self.assertEqual(len(while_op.inputs), orig_num_inputs + 2) # pylint: disable=g-deprecated-assert |
| 954 | |
| 955 | test() |
| 956 | |
| 957 | @test_util.run_deprecated_v1 |
| 958 | def testOpDef(self): |