Skip to content

Commit 6104f0d

Browse files
Dan Moldovantensorflower-gardener
Dan Moldovan
authored andcommitted
Strengthen input verification for SpecializeType by replacing DCHECK with explicit test/status return.
PiperOrigin-RevId: 453436708
1 parent c6c1755 commit 6104f0d

File tree

2 files changed

+18
-1
lines changed

2 files changed

+18
-1
lines changed

Diff for: tensorflow/core/framework/full_type_util.cc

+5-1
Original file line numberDiff line numberDiff line change
@@ -175,7 +175,11 @@ Status SubstituteVar(AttrMap& attrs, FullTypeDef& t) {
175175
}
176176

177177
Status SubstituteForEach(AttrMap& attrs, FullTypeDef& t) {
178-
DCHECK_EQ(t.args_size(), 3);
178+
if (t.args_size() != 3) {
179+
return Status(error::INVALID_ARGUMENT,
180+
absl::StrCat("illegal FOR_EACH type, expected 3 args, got ",
181+
t.args_size()));
182+
}
179183

180184
const auto& cont = t.args(0);
181185
const auto& tmpl = t.args(1);

Diff for: tensorflow/core/framework/full_type_util_test.cc

+13
Original file line numberDiff line numberDiff line change
@@ -510,6 +510,19 @@ TEST(SpecializeType, ForEachOverridesTargetOfNestedForEach) {
510510
EXPECT_EQ(t_actual.args(1).args(0).args(0).args_size(), 0);
511511
}
512512

513+
TEST(SpecializeType, ForEachRejectsMalformedInput) {
514+
OpDef op;
515+
FullTypeDef* t = op.add_output_arg()->mutable_experimental_full_type();
516+
t->set_type_id(TFT_FOR_EACH);
517+
t->add_args()->set_type_id(TFT_PRODUCT);
518+
519+
NodeDef ndef;
520+
AttrSlice attrs(ndef);
521+
522+
FullTypeDef ft;
523+
EXPECT_FALSE(SpecializeType(attrs, op, ft).ok());
524+
}
525+
513526
TEST(SpecializeType, RemovesLegacyVariant) {
514527
OpDef op;
515528
FullTypeDef* t = op.add_output_arg()->mutable_experimental_full_type();

0 commit comments

Comments
 (0)