private:
  std::unique_ptr<MLIRContext> context_;
  OpBuilder builder_;
  func::FuncOp func_;
  OwningOpRef<mlir::ModuleOp> module_;
};

Status MlirAbstractOp::Reset(const char* op, const char* device_name) {
  if (state_) {
    return FailedPrecondition("Reset called on already built op.");
  }
  TF_RETURN_IF_ERROR(
      tensorflow::OpRegistry::Global()->LookUpOpDef(op, &op_def_));
  assert(op_def_);

  for (int i = 0; i < num_devices; ++i) {
    const char* device_name = TF_DeviceListName(device_list, i, s);
    CHECK_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s);
    const char* device_type = TF_DeviceListType(device_list, i, s);
    CHECK_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s);
    LOG(INFO) << "Device " << i << " has name " << device_name << ", type "
              << device_type;
    if (string(device_type) == DEVICE_GPU) {