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

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";
}

  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>(

    "/tensorflow/core/tf2xla/v1/mlir_failed_xla_legalize_tf_count",
    "Counts the attempts of legalization of ops", "op_name");

auto *mlir_failed_legalization_count = tensorflow::monitoring::Counter<2>::New(
    "/tensorflow/core/tf2xla/v1/mlir_failed_xla_legalize_tf_pass_count",
    "Counts the failure of legalization of ops", "op_name", "legality");

class LegalizeTF : public impl::LegalizeTFBase<LegalizeTF> {
 public:

  bool IsAssertOrCheckNumerics(absl::string_view op_name) const {
    return op_name == "Assert" || op_name == "CheckNumerics";
  }

  bool IsStatefulRandomOp(absl::string_view op_name) const {
    return op_name == "RandomUniform" || op_name == "RandomShuffle" ||
           op_name == "RandomUniformInt" || op_name == "RandomStandardNormal" ||
           op_name == "TruncatedNormal" || op_name == "Multinomial";
  }

  string c_decl;
  for (int i = 0; i < op_info.arg_types.size(); ++i) {
    if (i > 0) strings::StrAppend(&c_decl, ", ");
    strings::StrAppend(&c_decl, op_info.arg_types[i], " ",
                       op_info.arg_names[i]);
  }
  if (include_attr && op_info.has_optional_attrs) {
    strings::StrAppend(&c_decl, ", const ", op_info.op_name, "::Attrs& attrs");
  }

  bool IsAllowListedOp(Operation *op) {
    std::string op_name = op->getName().getStringRef().str();
    if (IsAlwaysAllowlistedOp(op) || IsTfFakeQuantOp(op)) {
      return true;
    } else if (mode_ == kDefaultMode) {
      return QuantizableOpsInDefaultMode().count(op_name) > 0;
    } else if (mode_ == kLegacyIntegerMode) {
      return QuantizableOpsInLegacyMode().count(op_name) > 0;
    } else {

}
}  // namespace

Status GradientRegistry::Register(
    const string& op_name, GradientFunctionFactory gradient_function_factory) {
  auto iter = registry_.find(op_name);
  if (iter != registry_.end()) {
    const string error_msg = "Gradient already exists for op: " + op_name + ".";
    return errors::AlreadyExists(error_msg);
  }
  registry_.insert({op_name, gradient_function_factory});
  return absl::OkStatus();
}

def _find_op(
    graph: ops.Graph, op_name: Optional[str]
) -> Optional[ops.Operation]:
  """Finds the operation with `op_name`.

  Args:
    graph: The graph to find from.
    op_name: Name of the node.

  Returns:
    The operation that corresponds to `op_name`. Returns None iff op_name is an
    empty string or None.

  Raises:
    ValueError: `op_name` is malformed.
  """
  if not op_name: