diff --git a/tensorflow/compiler/xla/service/dynamic_parameter_binding.cc b/tensorflow/compiler/xla/service/dynamic_parameter_binding.cc index e9c8aa03e2a..7f0ae692f74 100644 --- a/tensorflow/compiler/xla/service/dynamic_parameter_binding.cc +++ b/tensorflow/compiler/xla/service/dynamic_parameter_binding.cc @@ -112,7 +112,8 @@ Status DynamicParameterBinding::Verify(const HloModule& module) const { return ForEachBinding([&](const DynamicParameter& dynamic_parameter, const DynamicDimension& dynamic_dimension) -> Status { - TF_RET_CHECK(dynamic_parameter.parameter_num < entry->num_parameters()); + TF_RET_CHECK(dynamic_parameter.parameter_num >= 0 && + dynamic_parameter.parameter_num < entry->num_parameters()); TF_RET_CHECK(dynamic_dimension.parameter_num < entry->num_parameters()); TF_RET_CHECK(ShapeUtil::IndexIsValid( entry->parameter_instruction(dynamic_parameter.parameter_num)->shape(), diff --git a/tensorflow/compiler/xla/service/hlo_parser.cc b/tensorflow/compiler/xla/service/hlo_parser.cc index 1a930370a9e..d7fa6b4a643 100644 --- a/tensorflow/compiler/xla/service/hlo_parser.cc +++ b/tensorflow/compiler/xla/service/hlo_parser.cc @@ -658,8 +658,14 @@ bool HloParser::ParseInstructionRhs(HloComputation::Builder* builder, tensorflow::int64 parameter_number; if (!ParseToken(TokKind::kLparen, "expects '(' before parameter number") || - !ParseInt64(¶meter_number) || - !ParseToken(TokKind::kRparen, "expects ')' after parameter number") || + !ParseInt64(¶meter_number)) { + return false; + } + if (parameter_number < 0) { + Error(lexer_.GetLoc(), "parameter number must be >= 0"); + return false; + } + if (!ParseToken(TokKind::kRparen, "expects ')' after parameter number") || !ParseAttributes(attrs)) { return false; } diff --git a/tensorflow/compiler/xla/service/hlo_parser_test.cc b/tensorflow/compiler/xla/service/hlo_parser_test.cc index ff52a89d1bd..6eee767900d 100644 --- a/tensorflow/compiler/xla/service/hlo_parser_test.cc +++ b/tensorflow/compiler/xla/service/hlo_parser_test.cc @@ -2558,5 +2558,13 @@ TEST_F(HloParserTest, ParseDynamicTuple) { << "actual: " << ShapeUtil::HumanString(actual); } +TEST_F(HloParserTest, NegativeParameterNumber) { + const string hlo_string = "par0 = f32[3,5] parameter(-1)"; + auto result = ParseHloString(hlo_string); + ASSERT_FALSE(result.status().ok()); + EXPECT_THAT(result.status().error_message(), + ::testing::HasSubstr("parameter number must be >= 0")); +} + } // namespace } // namespace xla