experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[XLA] Add support for parsing negative nans as constants in HLO parser.

Browse Source 
PiperOrigin-RevId: 335701644
Change-Id: Icedf22f31a99c8522e23e8ae95536a0cf1b7d767
This commit is contained in:
Marcello Maggioni 2020-10-06 12:56:52 -07:00 committed by TensorFlower Gardener
parent 527d4624fc
commit c05a9edacd
4 changed files with 26 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
tensorflow/compiler/xla/service/hlo_lexer.cc
View File

@ -387,6 +387,12 @@ TokKind HloLexer::LexNumberOrPattern() {
return TokKind::kNegInf;
}
static LazyRE2 neg_nan = {"-nan"};
if (RE2::Consume(&consumable, *neg_nan)) {
current_ptr_ = consumable.begin();
return TokKind::kNegNan;
}
return TokKind::kError;
}
@ -502,6 +508,8 @@ string TokKindToString(TokKind kind) {
return "kw_nan";
case TokKind::kw_inf:
return "kw_inf";
case TokKind::kNegNan:
return "kNegNan";
case TokKind::kNegInf:
return "kNegInf";
case TokKind::kPrimitiveType:

1
tensorflow/compiler/xla/service/hlo_lexer.h
View File

@ -65,6 +65,7 @@ enum class TokKind {
kw_nan,
kw_inf,
kNegNan, // -nan
kNegInf, // -inf
// Typed tokens.

4
tensorflow/compiler/xla/service/hlo_parser.cc
View File

@ -2717,6 +2717,7 @@ bool HloParserImpl::ParseDenseLiteral(Literal* literal, const Shape& shape) {
case TokKind::kInt:
case TokKind::kDecimal:
case TokKind::kw_nan:
case TokKind::kNegNan:
case TokKind::kw_inf:
case TokKind::kNegInf: {
add_one_elem_seen();
@ -4374,6 +4375,9 @@ bool HloParserImpl::ParseDouble(double* result) {
case TokKind::kw_nan:
*result = std::numeric_limits<double>::quiet_NaN();
break;
case TokKind::kNegNan:
*result = -std::numeric_limits<double>::quiet_NaN();
break;
case TokKind::kw_inf:
*result = std::numeric_limits<double>::infinity();
break;

13
tensorflow/compiler/xla/service/hlo_parser_test.cc
View File

@ -2120,6 +2120,19 @@ ENTRY %ShortConstant.v4 () -> f32[67,89] {
EXPECT_EQ(result.ValueOrDie()->ToString(HloPrintOptions()), original);
}
TEST_F(HloParserTest, NegativeNan) {
const string original = R"(HloModule NegativeNan_module
ENTRY %NegativeNan () -> bf16[2] {
ROOT %constant = bf16[2]{0} constant({-nan, -nan})
}
)";
auto result = ParseAndReturnUnverifiedModule(original);
EXPECT_EQ(Status::OK(), result.status());
EXPECT_EQ(result.ValueOrDie()->ToString(HloPrintOptions()), original);
}
TEST_F(HloParserTest, AttributesAnyOrder) {
const string original = R"(HloModule any_order_module