Function TEST_F

tensorflow/compiler/xla/tests/replay_test.cc:39–69

class ReplayTest : public ClientLibraryTestBase {};

TEST_F(ReplayTest, TwoPlusTwoReplay) {
  // Make 2+2 computation.
  XlaBuilder builder(TestName());
  auto two = ConstantR0<int32>(&builder, 2);
  Add(two, two);
  XlaComputation computation = builder.Build().ConsumeValueOrDie();

  // Serialize it out.
  std::unique_ptr<HloSnapshot> module =
      computation.Snapshot().ConsumeValueOrDie();

  // Replay it.
  XlaComputation replayed = client_->LoadSnapshot(*module).ConsumeValueOrDie();

  // Check signature is the same.
  std::unique_ptr<ProgramShape> original_shape =
      client_->GetComputationShape(computation).ConsumeValueOrDie();
  std::unique_ptr<ProgramShape> replayed_shape =
      client_->GetComputationShape(replayed).ConsumeValueOrDie();
  ASSERT_TRUE(protobuf_util::ProtobufEquals(original_shape->ToProto(),
                                            replayed_shape->ToProto()));

  // Run it.
  Literal literal =
      client_
          ->ExecuteAndTransfer(replayed, /*arguments=*/{}, &execution_options_)
          .ConsumeValueOrDie();

  // Expect 4.
  LiteralTestUtil::ExpectR0Equal<int32>(4, literal);
}

XLA_TEST_F(ReplayTest, XPlusYReplayWithParameters) {
  // Make computation.