TFE_Executor* executor = TFE_ContextGetExecutorForThread(ctx);
  TFE_ExecutorWaitForAllPendingNodes(executor, status);
  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);

  // TODO(b/136478427): Figure out how to correctly shut the server down.
  worker_server.release();

  // Update the server def with a new set of names (worker instead of
  // localhost).

          "Operation '", op.node.DebugString(), "' was changed by ",
          mutation_type,
          " after it was run by a session. This mutation will have no effect, "
          "and will trigger an error in the future. Either don't modify "
          "nodes after running them or create a new session.");
    }
  }
}

namespace {

// Helper method that creates a shape handle for a shape described by dims.

  ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
  TFE_DeleteExecutor(executor);
  TFE_DeleteContext(ctx);

  TF_DeleteStatus(status);

  // TODO(b/136478427): Figure out how to correctly shut the server down.
  worker_server.release();
}

TEST(CAPI, RemoteExecute) { TestRemoteExecute(false); }
TEST(CAPI, RemoteExecuteAsync) { TestRemoteExecute(true); }

  TFE_DeleteExecutor(executor);
  TFE_ContextRemoveFunction(ctx, "AddVariablesFunction", status);
  TFE_DeleteContext(ctx);

  TF_DeleteStatus(status);

  // TODO(b/136478427): Figure out how to correctly shut the server down.
  worker_server1.release();
  worker_server2.release();
}

TEST(CAPI, TestLocalFunctionWithPackedInput) {
  TestFunctionWithPackedInput(/*remote=*/false);
}

  ASSERT_EQ(errors::OK, status_.code()) << status_.message();

  ASSERT_NO_FATAL_FAILURE(CompareNumericalAndAutodiffGradients(
      ExpModel, BuildGradModel(ExpModel, registry_),
      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

        !DeviceNameUtils::IsSameAddressSpace(
            underlying_devices_[component_index], underlying_devices_[0])) {
      // Device names are from different address spaces, or we can't figure out
      // whether they are, so we'll fully-qualify everything.
      return underlying_devices_;
    }
  }
  std::vector<std::string> local_names;
  local_names.reserve(underlying_devices_.size());

  EXPECT_NE(TF_GetCode(status.get()), TF_CANCELLED);
  EXPECT_EQ(TF_GetCode(status.get()), TF_INVALID_ARGUMENT);
  EXPECT_THAT(TF_Message(status.get()), HasSubstr("assertion failed"));

  // Note that future collectives with the same context do not work at the
  // moment; once canceled, the collective executor requires the program to be
  // restarted / context to be reset.
}

        immediate_execution_ctx_.get(), X_vals, X_dims, 2, &X_raw);
    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
    X.reset(X_raw);
  }

  ASSERT_NO_FATAL_FAILURE(CompareNumericalAndAutodiffGradients(
      ReluModel, ReluGradModel, immediate_execution_ctx_.get(), {X.get()},
      UseFunction()));

  // Mathematically, Relu isn't differentiable at `0`. So `gradient_checker`

      ADD_FAILURE() << i << " -- " << TF_Message(status.get());
      continue;
    }
    auto tag = tensorflow::strings::StrCat("Device #", i, " (", name, ")");
    // Copy to device
    TFE_TensorHandle* hdevice =
        TFE_TensorHandleCopyToDevice(hcpu, ctx, name.c_str(), status.get());
    if (TF_GetCode(status.get()) != TF_OK) {
      ADD_FAILURE() << tag << " -- " << TF_Message(status.get());

  TFE_CancellationManager* c_mgr = TFE_NewCancellationManager();
  EXPECT_FALSE(TFE_CancellationManagerIsCancelled(c_mgr));

  TFE_CancelCallback callback1;
  callback1.callback = [](void* context) {
    ADD_FAILURE() << "Callback1 should be deregistered.";
  };
  TFE_CancellationToken token1 = TFE_CancellationManagerGetToken(c_mgr);
  EXPECT_TRUE(TFE_CancellationManagerRegisterCallback(c_mgr, token1, &callback1,