PatternRewriter& rewriter) const override {
    const auto f_attr = mlir::dyn_cast<FlatSymbolRefAttr>(op.getFAttr());
    // Non-quantizable op
    if (!op->hasAttr(kQuantTraitAttrName)) return failure();
    StringRef function_name = f_attr.getValue();
    // TODO(b/228928859): Improve the getter function to match attributes rather
    // than function name.

  // MLIR. This will also require some trait to identify optional attributes in
  // MLIR.
  constexpr char kShapeInferenceGraph[] = "shape_inference_graph";
  if (mlir::isa<mlir::TF::XlaHostComputeOp>(inst) &&
      !inst->hasAttr(kShapeInferenceGraph) &&
      !attrs_to_ignore.contains(kShapeInferenceGraph)) {
    AttrValue value;
    value.mutable_func()->set_name("");
    (*attributes)[kShapeInferenceGraph] = value;
  }

              quantizing_op, custom_map)) {
        if (!(enable_verify && enable_whole_model_verify)) {
          return failure();
        }
        if (quantizing_op->hasAttr(kDebugModeOpQuantAttrName) ||
            quantizing_op->hasAttr(kDebugModeOpFloatAttrName)) {
          return failure();
        }

        rewriter.setInsertionPoint(quantizing_op);
        Operation* float_op = rewriter.clone(*quantizing_op);

    func::FuncOp new_ref_func = dyn_cast<func::FuncOp>(ref_func->clone());
    return symbol_table.insert(new_ref_func);
  }

  LogicalResult match(LiftedOpT op) const override {
    if (!op->hasAttr(kQuantTraitAttrName) || op->getNumResults() != 1) {
      return failure();
    }

    Value result = op->getResult(0);
    for (auto user : result.getUsers()) {

  auto tf_dialect = ctx->getLoadedDialect<TF::TensorFlowDialect>();
  func.walk([&](Operation *op) {
    if (op->getDialect() != tf_dialect) return;
    if (IsAllowListedOp(op)) return;
    if (op->hasAttr(kNoFallbackAttr)) {
      op->removeAttr(kNoFallbackAttr);
      return;
    }
    if (!ConvertToFlexOp(op)) signalPassFailure();
  });
}
}  // namespace internal

                              const int operand_index) {
  auto* defining_op = op->getOperand(operand_index).getDefiningOp();
  for (auto attr : kQuantizationAxisAttrs) {
    if (defining_op->hasAttr(attr)) {
      return defining_op->getAttr(attr);
    }
  }
  // Not found.
  return rewriter.getI64IntegerAttr(-1);
}

LogicalResult CheckIfAttrIs8Bit(const std::string& attr, Operation* op,

    if (isa<TPUPartitionedCallOp>(call)) return false;
    if (isa<RemoteCallOp>(call)) return false;
    // Maintain inlining for  `tf.function`s with jit_compile option.
    if (callable->hasAttr("tf._XlaMustCompile")) return true;
    auto noinline_attr_name = absl::StrCat("tf.", tensorflow::kNoInlineAttr);
    if (auto noinline_attr =
            callable->getAttrOfType<BoolAttr>(noinline_attr_name))

      op->getAttrOfType<ArrayAttr>(kTfSavedModelExportedNamesAttr);
  return exported_names && !exported_names.empty();
}

// Check if a function is an entry function.
bool IsEntryFunction(func::FuncOp op) {
  return op->hasAttr(kEntryFunctionAttr);
}

// Returns true iff the provided FuncOp is qualified to be included in the main
// function.
bool ShouldIncludeInMainFunction(func::FuncOp func_op) {

bool IsQuantizedXlaCallModuleOp(TF::XlaCallModuleOp xla_call_module_op) {
  return !IsOpNotQuantized(xla_call_module_op) &&
         xla_call_module_op->hasAttr(kQuantTraitAttrName) &&
         xla_call_module_op->hasAttr(kEntryFuncAttrName);
}

// Returns the entry function, i.e. the callee of `xla_call_module_op`.
func::FuncOp GetEntryFuncOp(TF::XlaCallModuleOp xla_call_module_op,

    }

    if (op.hasAttr(mlir::TF::kReplicationInfoAttr)) {
      // For replicated case, borrow cluster structure from replication info.
      // Following condition is already checked in
      // `HasValidCompilationAndReplicationAttributes` above, assert here for
      // documentation and to avoid breakage when that function is changed.
      assert(op.hasAttr(mlir::TF::kCompileDeviceTypeAttr));