Strengthen input verification for SpecializeType by replacing DCHECK …

…with explicit test/status return. PiperOrigin-RevId: 453436708
tensorflow · Jun 7, 2022 · 6104f0d · 6104f0d
1 parent c6c1755
commit 6104f0d
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 1 deletion.
diff --git a/tensorflow/core/framework/full_type_util.cc b/tensorflow/core/framework/full_type_util.cc
@@ -175,7 +175,11 @@ Status SubstituteVar(AttrMap& attrs, FullTypeDef& t) {
 }
 
 Status SubstituteForEach(AttrMap& attrs, FullTypeDef& t) {
-  DCHECK_EQ(t.args_size(), 3);
+  if (t.args_size() != 3) {
+    return Status(error::INVALID_ARGUMENT,
+                  absl::StrCat("illegal FOR_EACH type, expected 3 args, got ",
+                               t.args_size()));
+  }
 
   const auto& cont = t.args(0);
   const auto& tmpl = t.args(1);

diff --git a/tensorflow/core/framework/full_type_util_test.cc b/tensorflow/core/framework/full_type_util_test.cc
@@ -510,6 +510,19 @@ TEST(SpecializeType, ForEachOverridesTargetOfNestedForEach) {
   EXPECT_EQ(t_actual.args(1).args(0).args(0).args_size(), 0);
 }
 
+TEST(SpecializeType, ForEachRejectsMalformedInput) {
+  OpDef op;
+  FullTypeDef* t = op.add_output_arg()->mutable_experimental_full_type();
+  t->set_type_id(TFT_FOR_EACH);
+  t->add_args()->set_type_id(TFT_PRODUCT);
+
+  NodeDef ndef;
+  AttrSlice attrs(ndef);
+
+  FullTypeDef ft;
+  EXPECT_FALSE(SpecializeType(attrs, op, ft).ok());
+}
+
 TEST(SpecializeType, RemovesLegacyVariant) {
   OpDef op;
   FullTypeDef* t = op.add_output_arg()->mutable_experimental_full_type();