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

  const char remote_device_name[] =
      "/job:localhost/replica:0/task:1/device:CPU:0";
  const char local_device_name[] =
      "/job:localhost/replica:0/task:0/device:CPU:0";
  CheckRemoteMatMulExecutesOK(ctx, remote_device_name, local_device_name);

  TFE_Executor* executor = TFE_ContextGetExecutorForThread(ctx);

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

  const char task0_name[] = "/job:localhost/replica:0/task:0/device:CPU:0";
  const char task1_name[] = "/job:localhost/replica:0/task:1/device:CPU:0";
  const char task2_name[] = "/job:localhost/replica:0/task:2/device:CPU:0";

  // Create one variable per task.
  TFE_TensorHandle* h0 = TestVariable(ctx, 1.0, task1_name);

  }
  auto* config = server_def.mutable_default_session_config();
  config->mutable_experimental()->set_collective_group_leader(
      tensorflow::strings::StrCat("/job:", job_name, "/replica:0/task:", 0));
  auto* rewrite_options =
      config->mutable_graph_options()->mutable_rewrite_options();
  rewrite_options->set_scoped_allocator_optimization(
      tensorflow::RewriterConfig::ON);

                              const char* prefix, TF_Status* status) {
  std::vector<std::unique_ptr<tensorflow::Device>> devices;

  if (prefix == nullptr || strlen(prefix) == 0)
    prefix = "/job:localhost/replica:0/task:0";

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

    return nullptr;
  }
  std::vector<std::unique_ptr<tensorflow::Device>> devices;
  status->status = tensorflow::DeviceFactory::AddDevices(
      opts->session_options.options, "/job:localhost/replica:0/task:0",
      &devices);
  if (!status->status.ok()) return nullptr;
  std::unique_ptr<tensorflow::DeviceMgr> device_mgr(
      new tensorflow::DynamicDeviceMgr(std::move(devices)));

      CreateContext(serialized_server_def_1,
                    /*isolate_session_state=*/false, init_timeout_in_ms);

  // Remote device on `worker2`.
  const char remote_device[] = "/job:localhost/replica:0/task:2/device:CPU:0";
  // `ctx_0`, `ctx_1` contains `remote_device`.
  {
    const std::vector<std::string>& device_names = ListDeviceNames(ctx_0);
    ASSERT_TRUE(std::find(device_names.begin(), device_names.end(),