context.get(), components, device_name, status.get());
  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());

  // Run a collective sum, so each component should now be the same.
  TensorHandlePtr reduced(
      CollectiveSum(context.get(), parallel_value.get(), 2, status.get()));
  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());

  // Now load a block of a different file "b" at timestamp `now` + 1
  env->SetNowSeconds(now + 1);
  TF_EXPECT_OK(ReadCache(&cache, "b", 0, 1, &out));
  // Now load a different block of file "a" at timestamp `now` + 1. When the
  // first block of "a" expires, this block should also be removed because it
  // also belongs to file "a".
  TF_EXPECT_OK(ReadCache(&cache, "a", 8, 1, &out));

  while (!stop_pruning_thread_.WaitForNotificationWithTimeout(
      absl::Microseconds(1000000))) {
    absl::MutexLock lock(&mu_);
    uint64_t now = timer_seconds_();
    while (!lra_list_.empty()) {
      auto it = block_map_.find(lra_list_.back());
      if (now - it->second->timestamp <= max_staleness_) {
        // The oldest block is not yet expired. Come back later.
        break;
      }

  TF_DeleteAbstractOp(matmul_op);

  TF_AbstractFunction* func =
      TF_FinalizeFunction(graph_ctx, mm_outputs, status.get());
  ASSERT_EQ(TF_OK, TF_GetCode(status.get())) << TF_Message(status.get());

  /* Now that the graph is built, test graph implementation on matmul example:
    [[1,1] ,   *   [[1,1] ,   =  [[2,2],
     [1,1]]         [1,1]]        [2,2]]
  */

  // Build eager context.