// eager runtime.
  if (enable_grappler) {
    tensorflow::ConfigProto config;
    // Do not skip grappler optimization even for small graphs.
    config.mutable_graph_options()
        ->mutable_rewrite_options()
        ->set_min_graph_nodes(-1);
    string serialized_config;
    ASSERT_TRUE(config.SerializeToString(&serialized_config));
    TFE_OpSetAttrString(
        op, "config_proto",

                                   PJRT_NamedValue* create_options,
                                   int num_options) {
  absl::StatusOr<std::unique_ptr<xla::PjRtClient>> pjrt_client =
      xla::GetCApiClient(device_type, pjrt::ConvertFromPjRtNamedValueList(
                                          create_options, num_options));
  if (!pjrt_client.ok()) {
    status->status = pjrt_client.status();
    return;

class GraphErrorInjectionPass : public tensorflow::GraphOptimizationPass {
 public:
  static bool enabled_;
  GraphErrorInjectionPass() {}

  tensorflow::Status Run(
      const tensorflow::GraphOptimizationPassOptions& options) override {
    if (!enabled_) {
      return absl::OkStatus();
    }
    if (first_call_) {
      first_call_ = false;
      return absl::OkStatus();
    }

// integration testing rather than purely testing the parallel device. They
// correspond fairly well to the implementation, but testing the C++ directly is
// another option.

namespace tensorflow {
namespace parallel_device {

Variable* Variable::Create(TFE_Context* context, TF_DataType type,
                           const int64_t* dims, const int num_dims,

// integration testing rather than purely testing the parallel device. They
// correspond fairly well to the implementation, but testing the C++ directly is
// another option.

namespace tensorflow {
namespace parallel_device {

using ::testing::HasSubstr;

TEST(PARALLEL_DEVICE, TestBasicCPU) {
  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(

      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);
  rewrite_options->mutable_scoped_allocator_opts()->add_enable_op(
      "CollectiveReduce");

void TFE_ContextOptionsSetConfig(TFE_ContextOptions* options, const void* proto,
                                 size_t proto_len, TF_Status* status) {
  TF_SetConfig(&options->session_options, proto, proto_len, status);
}

void TFE_ContextOptionsSetAsync(TFE_ContextOptions* options,
                                unsigned char enable) {
  options->async = enable;
}

void TFE_ContextOptionsSetDevicePlacementPolicy(

 protected:
  ShapeInferenceTest()
      : status_(TF_NewStatus()), tfe_context_options_(TFE_NewContextOptions()) {
    tfe_context_ = TFE_NewContext(tfe_context_options_, status_);
    CHECK_EQ(TF_OK, TF_GetCode(status_)) << TF_Message(status_);
  }

  ~ShapeInferenceTest() override {
    TFE_DeleteContextOptions(tfe_context_options_);
    TFE_DeleteContext(tfe_context_);
    TF_DeleteStatus(status_);
  }

  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);

  // Remove the device that has the host device name since host device is alreay
  // in an initialized context.