| 69 | } |
| 70 | |
| 71 | XLA_TEST_F(ReplayTest, XPlusYReplayWithParameters) { |
| 72 | // Make computation. |
| 73 | XlaBuilder builder(TestName()); |
| 74 | auto x = Parameter(&builder, 0, ShapeUtil::MakeShape(S32, {}), "x"); |
| 75 | auto y = Parameter(&builder, 1, ShapeUtil::MakeShape(S32, {}), "y"); |
| 76 | Add(x, y); |
| 77 | XlaComputation computation = builder.Build().ConsumeValueOrDie(); |
| 78 | |
| 79 | // Serialize it out. |
| 80 | std::unique_ptr<HloSnapshot> module = |
| 81 | computation.Snapshot().ConsumeValueOrDie(); |
| 82 | |
| 83 | // Replay it. |
| 84 | XlaComputation replayed = client_->LoadSnapshot(*module).ConsumeValueOrDie(); |
| 85 | |
| 86 | // Check signature is the same. |
| 87 | std::unique_ptr<ProgramShape> original_shape = |
| 88 | client_->GetComputationShape(computation).ConsumeValueOrDie(); |
| 89 | std::unique_ptr<ProgramShape> replayed_shape = |
| 90 | client_->GetComputationShape(replayed).ConsumeValueOrDie(); |
| 91 | ASSERT_TRUE(protobuf_util::ProtobufEquals(original_shape->ToProto(), |
| 92 | replayed_shape->ToProto())); |
| 93 | |
| 94 | // Run it. |
| 95 | std::unique_ptr<GlobalData> x_data = |
| 96 | client_->TransferToServer(LiteralUtil::CreateR0<int32>(2)) |
| 97 | .ConsumeValueOrDie(); |
| 98 | std::unique_ptr<GlobalData> y_data = |
| 99 | client_->TransferToServer(LiteralUtil::CreateR0<int32>(3)) |
| 100 | .ConsumeValueOrDie(); |
| 101 | Literal literal = |
| 102 | client_ |
| 103 | ->ExecuteAndTransfer(replayed, |
| 104 | /*arguments=*/{x_data.get(), y_data.get()}, |
| 105 | &execution_options_) |
| 106 | .ConsumeValueOrDie(); |
| 107 | |
| 108 | // Expect 5. |
| 109 | LiteralTestUtil::ExpectR0Equal<int32>(5, literal); |
| 110 | } |
| 111 | |
| 112 | TEST_F(ReplayTest, MapPlusTwoOverR1) { |
| 113 | // As above, but with map(+2) over some constant array. |
nothing calls this directly
no test coverage detected