Merge pull request #129186 from bodapatisaikrishna:fix/62982-loopcond-empty-frame

PiperOrigin-RevId: 1005382776

tensorflow/core/graph/graph_partition.cc

@@ -24,6 +24,8 @@ limitations under the License.
 #include <vector>
 
 #include "absl/container/flat_hash_map.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
 #include "tensorflow/core/framework/function.h"
 #include "tensorflow/core/framework/memory_types.h"
 #include "tensorflow/core/framework/node_def.pb.h"
@@ -36,6 +38,7 @@ limitations under the License.
 #include "tensorflow/core/graph/graph.h"
 #include "tensorflow/core/graph/graph_debug_info_builder.h"
 #include "tensorflow/core/graph/graph_def_builder.h"
+#include "tensorflow/core/graph/graph_node_util.h"
 #include "tensorflow/core/graph/node_builder.h"
 #include "tensorflow/core/graph/tensor_id.h"
 #include "tensorflow/core/lib/core/errors.h"
@@ -646,7 +649,10 @@ absl::Status AddControlFlow(const PartitionOptions& opts, Graph* g,
 
     if (IsLoopCond(node)) {
       const std::string& frame_name = cf_info[node->id()].frame_name;
-      DCHECK(!frame_name.empty());
+      if (frame_name.empty()) {
+        return absl::InvalidArgumentError(absl::StrCat(
+            FormatNodeForError(*node), " must be inside a while loop frame."));
+      }
       frame_cond_map[frame_name] = node;
     }
   }

tensorflow/core/graph/graph_partition_test.cc

@@ -22,6 +22,7 @@ limitations under the License.
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/strings/match.h"
 #include "absl/strings/str_cat.h"
 #include "tensorflow/cc/ops/array_ops.h"
 #include "tensorflow/cc/ops/const_op.h"
@@ -487,6 +488,32 @@ TEST_F(GraphPartitionTest, PartitionIncompleteGraph) {
   EXPECT_EQ(error::INVALID_ARGUMENT, status.code()) << status;
 }
 
+TEST_F(GraphPartitionTest, LoopCondOutsideWhileLoop) {
+  auto a1 = BoolInput(in_.WithOpName("A1"));
+  LoopCond(in_.WithOpName("A2"), a1);
+
+  Graph g(OpRegistry::Global());
+  TF_ASSERT_OK(
+      ConvertGraphDefToGraph(GraphConstructorOptions(), ToGraphDef(), &g));
+  for (Node* node : g.nodes()) {
+    node->set_assigned_device_name(DeviceName(node));
+  }
+
+  PartitionOptions popts;
+  popts.node_to_loc = SplitByDevice;
+  popts.new_name = [&g](const std::string& prefix) {
+    return g.NewName(prefix);
+  };
+  popts.get_incarnation = [](const std::string&) { return 1; };
+
+  std::unordered_map<std::string, GraphDef> partitions;
+  absl::Status status = Partition(popts, &g, &partitions);
+  EXPECT_EQ(error::INVALID_ARGUMENT, status.code()) << status;
+  EXPECT_TRUE(
+      absl::StrContains(status.message(), "must be inside a while loop frame"))
+      << status;
+}
+
 TEST_F(GraphPartitionTest, Functions) {
   FunctionDefLibrary fdef_lib;
   *fdef_lib.add_function() = test::function::XTimesTwo();

tensorflow/python/kernel_tests/control_flow/control_flow_ops_py_test.py

@@ -4838,6 +4838,17 @@ class ControlFlowContextCheckTest(test.TestCase):
       tf_cond.cond(
           math_ops.less(1, 2), true_fn, lambda: constant_op.constant(0))
 
+  @test_util.run_deprecated_v1
+  def testLoopCondOutsideWhileLoop(self):
+    # LoopCond outside a while loop frame must fail with InvalidArgumentError
+    # rather than aborting during graph partitioning (issue #62982).
+    with self.cached_session() as sess:
+      cond = gen_control_flow_ops.loop_cond(input=True)
+      with self.assertRaisesRegex(
+          errors_impl.InvalidArgumentError, "must be inside a while loop frame"
+      ):
+        sess.run(cond)
+
 
 class TupleTest(test.TestCase):