namespace odml {

static bool isDenseI64Array(llvm::StringRef op_name,
                            llvm::StringRef field_name) {
  if (op_name == "stablehlo.broadcast" && field_name == "broadcast_sizes")
    return true;
  if (op_name == "stablehlo.broadcast_in_dim" &&
      field_name == "broadcast_dimensions")
    return true;
  if ((op_name == "stablehlo.convolution" ||
       op_name == "stablehlo.dynamic_conv") &&

          // Conditions when using TOCO.
          if (absl::StrContains(op_name, "LSTMOp")) continue;
        } else {
          // Conditions when using MLIR.
          if (!(absl::StrContains(op_name, "LSTMOp") ||
                absl::StrContains(op_name, "SoftmaxOp") ||
                absl::StrContains(op_name, "LogisticOp") ||
                absl::StrContains(op_name, "L2NormalizationOp") ||

  attrs_[attr_name] = TypeAttr::get(mlir_type);
  return absl::OkStatus();
}

Status MlirAbstractOp::SetOpName(const char* const op_name) {
  // TODO(aminim): should we use a location?
  if (op_name_) {
    return FailedPrecondition("SetOpName called on already built op.");
  }
  op_name_ = op_name;
  return absl::OkStatus();
}

Status MlirAbstractOp::AddRef(Type type, Type* output_type) {

  // Find the op names with tf or tfl dialect prefix, Ex: "tf.Abs" or "tfl.Abs".
  auto collectOps = [this](Operation *op) {
    const auto &op_name = op->getName().getStringRef().str();
    if (absl::StartsWith(op_name, "tf.") || absl::StartsWith(op_name, "tfl.")) {
      loc_to_name_.emplace(op->getLoc(), op_name);
    }
  };

  for (auto &region : module->getRegions()) {
    region.walk(collectOps);
  }

static inline std::string GetOperatorBuilderName(StringRef op_name) {
  assert(op_name.starts_with("TFL_") && "unexpected op prefix");
  assert(op_name.ends_with("Op") && "unexpected op suffix");

  // E.g., AddOp -> CreateAddOperator
  std::ostringstream oss;
  oss << "Create" << op_name.drop_front(4).str() << "erator";
  return oss.str();
}

static inline bool IsLstmOp(const StringRef op_name) {
  return op_name.take_back(6) == "LSTMOp";
}

  const char* node_name = "SomeNodeName";
  const char* op_name = "FooOp";
  const char* device_name = "FakeDeviceName1";

  REGISTER_OP(op_name)
      .Input("input1: double")
      .Input("input2: uint8")
      .Output("output1: uint8")
      .Attr("SomeDataTypeAttr: type");

  TF_KernelBuilder* builder = TF_NewKernelBuilder(
      op_name, device_name, &MyCreateFunc, &MyComputeFunc, &MyDeleteFunc);

  {

  if (!op_name_or_status.ok()) return false;

  auto op_name = std::move(op_name_or_status).value();
  if (op_name == "NextIteration")
    return mlir::isa<mlir::TF::TensorFlowRefType>(
        mlir::getElementTypeOrSelf(op->getOperand(0).getType()));

  if (op_name == "Enter" || op_name == "Exit" || op_name == "Switch" ||
      op_name == "Merge") {
    return mlir::isa<mlir::TF::TensorFlowRefType>(

namespace tracing {

Status MaybeSetOpName(AbstractOperation* op, const char* op_name) {
  if (isa<TracingOperation>(op)) {
    TF_RETURN_IF_ERROR(dyn_cast<TracingOperation>(op)->SetOpName(op_name));
  }
  if (isa<gradients::TapeOperation>(op)) {
    TF_RETURN_IF_ERROR(MaybeSetOpName(
        dyn_cast<gradients::TapeOperation>(op)->GetBackingOperation(),
        op_name));
  }
  return absl::OkStatus();
}
}  // namespace tracing

using tensorflow::OpRegistrationData;
using tensorflow::OpRegistry;
using tensorflow::Status;

namespace {

void Register(const std::string& op_name, OpRegistry* registry) {
  registry->Register([op_name](OpRegistrationData* op_reg_data) -> Status {
    op_reg_data->op_def.set_name(op_name);
    return absl::OkStatus();
  });
}

}  // namespace

TEST(TfTextUtilsTest, TestTfTextRegistered) {

}

}  // namespace

void RegisterCostFunction(absl::string_view op_name,
                          CostFunction cost_function) {
  RegisterCostFunction(GetCostFunctionRegistry(), op_name,
                       std::move(cost_function));
}

bool HasCostFunctionRegistered(absl::string_view op_name) {
  return GetCostFunctionRegistry().contains(op_name);
}

int64_t CostAnalysis::GetCost(mlir::Operation* op) const {