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

Add HLO test with tuple-shaped output containing a token.

Browse Source 
Also, fix CHECK failure in proto deserialization when the shape of an infeed was not a tuple.

PiperOrigin-RevId: 221209419
This commit is contained in:
Mark Heffernan 2018-11-12 21:35:58 -08:00 committed by TensorFlower Gardener
parent 54b65b83a3
commit b32d3167ff
3 changed files with 21 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tensorflow/compiler/xla/literal_comparison.cc
View File

@ -141,8 +141,10 @@ int64 RecursiveElementCount(const Shape& shape) {
total += RecursiveElementCount(ShapeUtil::GetTupleElementShape(shape, i)); total += RecursiveElementCount(ShapeUtil::GetTupleElementShape(shape, i));
} }
return total; return total;
} else { } else if (ShapeUtil::IsArray(shape)) {
return ShapeUtil::ElementsIn(shape); return ShapeUtil::ElementsIn(shape);
} else {
return 0;
} }
} }

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

@ -312,6 +312,10 @@ StatusOr<std::unique_ptr<HloInstruction>> HloInstruction::CreateFromProto(
proto.exponent_bits(), proto.mantissa_bits()); proto.exponent_bits(), proto.mantissa_bits());
break; break;
case HloOpcode::kInfeed: { case HloOpcode::kInfeed: {
TF_RET_CHECK(ShapeUtil::IsTuple(proto.shape()) &&
(ShapeUtil::TupleElementCount(proto.shape()) == 2))
<< "Infeed should have a tuple shape with 2 operands, but has: "
<< proto.shape();
const Shape& data_shape = const Shape& data_shape =
ShapeUtil::GetTupleElementShape(proto.shape(), 0); ShapeUtil::GetTupleElementShape(proto.shape(), 0);
TF_RET_CHECK(proto.operand_ids_size() == 1) TF_RET_CHECK(proto.operand_ids_size() == 1)

14
tensorflow/compiler/xla/tests/token_hlo_test.cc
View File

@ -38,6 +38,20 @@ XLA_TEST_F(TokenHloTest, SingleTokenInstruction) {
EXPECT_TRUE(LiteralTestUtil::Equal(result, LiteralUtil::CreateToken())); EXPECT_TRUE(LiteralTestUtil::Equal(result, LiteralUtil::CreateToken()));
} }
XLA_TEST_F(TokenHloTest, TokenInTuple) {
std::unique_ptr<HloModule> module = CreateNewUnverifiedModule();
auto builder = HloComputation::Builder(TestName());
auto token = builder.AddInstruction(HloInstruction::CreateToken());
builder.AddInstruction(HloInstruction::CreateTuple({token}));
module->AddEntryComputation(builder.Build());
TF_ASSERT_OK_AND_ASSIGN(Literal result, Execute(std::move(module), {}));
Literal token_literal = LiteralUtil::CreateToken();
EXPECT_TRUE(
LiteralTestUtil::Equal(result, LiteralUtil::MakeTuple({&token_literal})));
}
XLA_TEST_F(TokenHloTest, TokenTree) { XLA_TEST_F(TokenHloTest, TokenTree) {
std::unique_ptr<HloModule> module = CreateNewUnverifiedModule(); std::unique_ptr<HloModule> module = CreateNewUnverifiedModule();
auto builder = HloComputation::Builder(TestName()); auto builder = HloComputation::Builder(TestName());