TFE_OpSetDevice(matmul, cpu_device_name.c_str(), status);
    EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
    auto remote_arg =
        tensorflow::TensorHandleFromInterface(tensorflow::unwrap(h1_task2));
    // The input handles should never change since they have been mirrored.
    ASSERT_FALSE(remote_arg->HasLocalMirror(nullptr));
  }

  TFE_TensorHandle* retvals[1];
  int num_retvals = 1;

  TFE_TensorHandle* var_handle0 = TestVariable(ctx, 1.0, dev0_name);
  EXPECT_NE(var_handle0, nullptr);

  tensorflow::Status status2;
  EXPECT_EQ(tensorflow::unwrap(var_handle0)->DeviceName(&status2), dev0_name);

  // Rename local device
  // This server def has the task index set to 0.
  string serialized = server_def.SerializeAsString();

    return;
  }
  status->status = tensorflow::unwrap(ctx)->EnableCollectiveOps(server_def);
}

TF_CAPI_EXPORT extern void TFE_AbortCollectiveOps(TFE_Context* ctx,
                                                  TF_Status* status) {
  tensorflow::EagerContext* context =
      tensorflow::ContextFromInterface(tensorflow::unwrap(ctx));
  auto collective_executor_handle = context->GetCollectiveExecutorHandle();

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

  // Test that full type information can be accessed.
  auto outs = unwrap(add_outputs);
  auto h = outs->outputs[0];
  ASSERT_NE(h, nullptr);
  ASSERT_EQ(h->FullType().type_id(), TFT_UNSET);
  ASSERT_EQ(unwrap(inputs[0])->FullType().type_id(), TFT_UNSET);

  // Clean up operation and inputs.
  TF_DeleteAbstractOp(add_op);

  int real_num_outputs = expected_max_outputs;
  TFE_OpSetCancellationManager(op_.get(), wrap(cancellation_manager_), status);
  if (TF_GetCode(status) != TF_OK) return;

  // unwrap op_ and set step_id only if valid step id value was set.
  // Currently only required for non-TFRT use cases, e.g., EagerOp.
  if (step_id_.has_value()) {
    tensorflow::unwrap(op_.get())->SetStepId(step_id_.value());
  }

      } else {
        TF_SetStatus(
            status, TF_INVALID_ARGUMENT,
            absl::StrCat(
                "Got a non-parallel tensor ",
                tensorflow::unwrap(absl::get<TFE_TensorHandle*>(input))
                    ->DebugString(),
                " as input to a parallel operation. First pack non-parallel "
                "tensors for each device into a parallel tensor explicitly.")

            {"merge_op", "Add"},
            {"final_op", "Id"},
            {"subdiv_offsets", std::vector<int>()}},
           /*dep=*/{"assert"}},
      });
  TF_ASSERT_OK(ContextFromInterface(unwrap(context.get()))
                   ->AddFunctionDef(assert_and_collective));

  std::unique_ptr<TFE_Op, decltype(&TFE_DeleteOp)> call_op(
      TFE_NewOp(context.get(), "AssertAndCollective", status.get()),

        "//tensorflow/c/eager:abstract_operation",
        "//tensorflow/c/eager:abstract_tensor_handle",
        "//tensorflow/c/eager:c_api_unified_internal",
    ],
)

filegroup(
    name = "pywrap_required_hdrs",
    srcs = [
        "array_ops.h",
        "io_ops.h",
        "math_ops.h",
        "nn_ops.h",
        "resource_variable_ops.h",
    ],

)

# Currently pybind extension shared objects must use only C API headers since
# the C API has static initializers duplicated in the Python bindings. So we
# need a second rule that omits .cc files, in
# tensorflow/python:_pywrap_parallel_device.
filegroup(
    name = "lib_headers",
    srcs = ["parallel_device_lib.h"],
)

filegroup(
    name = "lib_sources",
    srcs = ["parallel_device_lib.cc"],
)

filegroup(

      context.get(), components, device_name, status.get());
  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());
  ImmediateExecutionTensorHandle* unwrapped_handle =
      tensorflow::unwrap(device_value.get());
  std::string summarized;
  TF_ASSERT_OK(unwrapped_handle->SummarizeValue(summarized));
  EXPECT_THAT(summarized, HasSubstr("\"CPU:0\": 3"));
}

}  // namespace parallel_device