immediate_execution_ctx_.get(), {x.get()}, UseFunction()));
}

TEST_P(CppGradients, TestMatMulGrad) {
  // TODO(vnvo2409): Figure out why `gradient_checker` does not work very
  // well with `MatMul` and remove `TestMatMul*` in
  // `mnist_gradients_test` when done.
  GTEST_SKIP();

  float A_vals[] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f};
  int64_t A_dims[] = {3, 3};
  AbstractTensorHandlePtr A;
  {

  // Check that copying it to the old remote device (named localhost) fails.
  TFE_TensorHandleCopyToDevice(h0_task0_new, ctx, remote_device_name, status);
  EXPECT_NE(TF_OK, TF_GetCode(status)) << TF_Message(status);

  // Copying and executing on the new remote device works.
  const char new_remote_device_name[] =
      "/job:worker/replica:0/task:1/device:CPU:0";

  // Set device to CPU since there is no version of split for int32 on GPU
  // TODO(iga): Convert all these helpers and tests to use floats because
  // they are usually available on GPUs. After doing this, remove TF_SetDevice
  // call in c_api_function_test.cc
  TF_SetDevice(desc, "/cpu:0");
  *op = TF_FinishOperation(desc, s);
  ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s);
  ASSERT_NE(*op, nullptr);
}

  // and not for correctness. The only downside of keeping this 1 seems to be
  // that the gradient accumulation is unbounded and we will never
  // aggressively aggregate accumulated gradients to recover memory.
  // Revisit and fix.
  return 1;
}

// Consumes references to the tensors in the gradient_tensors list and returns
// a tensor with the result.
AbstractTensorHandle* TapeVSpace::AggregateGradients(

  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);

  // Pack 3 variable handles into one TFE_TensorHandle.
  // When remote is false, function device is placed on task0. Handle types are
  // REMOTE, REMOTE, LOCAL on task0. When remote is true, function device is
  // placed on task1, Handle types are LOCAL, REMOTE, LOCAL on task1.
  int num_replicas = 3;
  std::vector<TFE_TensorHandle*> handles = {h0, h1, h2};

    const TFE_OpAttrs* attributes, int expected_max_outputs,
    TF_Status* status) {
  absl::optional<std::vector<MaybeParallelTensorOwned>> result;
  // TODO(allenl): We should remove "TPU" from these op names at the very least,
  // or consider other ways of packing/unpacking parallel tensors.
  if (operation_name == std::string("TPUReplicatedInput")) {

    per_device_output_tensors.push_back(device_thread->Join(status));
    // We will run every Join even if there are bad statuses in case the user
    // wants to recover and continue running ops on the parallel device (which
    // would otherwise deadlock).
    if (TF_GetCode(status) != TF_OK &&
        (first_bad_status == nullptr

  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));
  // Ensure that all blocks are in the cache (i.e. reads are cache hits).
  EXPECT_EQ(cache.CacheSize(), 24);

}

TEST_F(CApiExperimentalFunctionTest, EmptyGraphRemoveNonExistentFunction) {
  TF_GraphRemoveFunction(func_graph_, "wrong_name", s_);
  EXPECT_EQ(TF_INVALID_ARGUMENT, TF_GetCode(s_));
  EXPECT_EQ(string("Tried to remove non-existent function 'wrong_name'."),
            string(TF_Message(s_)));
}

TEST_F(CApiExperimentalFunctionTest, GraphRemoveNonExistentFunction) {
  TF_Function* funcs[1];
  DefineFunction(func_name_, &func_);

  tensorflow::SessionOptions sess_options;
  (*sess_options.config.mutable_device_count())["CPU"] = num_cpus;
  status->status =
      tensorflow::DeviceFactory::AddCpuDevices(sess_options, prefix, &devices);

  // Remove the device that has the host device name since host device is alreay
  // in an initialized context.
  for (auto d = devices.begin(); d != devices.end();) {
    if (absl::StrContains(d->get()->name(), "CPU:0")) {