if (op->hasAttr(TF::kReplicationInfoAttr)) {
    op->setAttr(TF::kReplicationInfoAttr, attr);
  }
}

// Replaces the replication region attribute if it already exists.
void UpdateReplicationAttr(TF::TPUCompilationResultOp& op, StringAttr attr) {
  // Special case for getting the replication region for
  // TPUCompilationResultsOp.
  if (op->hasAttr(kTpuCompilationStatus)) {

  // A map to find an attribute from its identifier.
  llvm::StringMap<Attribute> identifier_to_attr;

  for (Operation& inner_op : float_func.getBody().front().getOperations()) {
    if (!inner_op.hasAttr(kAttrMapAttribute)) continue;
    // Insert quantization related attribute if they exists. Quantization
    // attributes are generated in the prepare pass so the attr_map doesn't
    // contain the attribute names.

        if (walked_op->hasAttr(kXlaOutsideCompilationAttr)) {
          return WalkResult::interrupt();
        }
        return WalkResult::advance();
      })
      .wasInterrupted();
}

// Returns whether `op` or any ancestors of `op` are outside compiled.
bool HasOutsideCompilationAncestor(Operation* op) {
  while (op) {
    if (op->hasAttr(kXlaOutsideCompilationAttr)) {
      return true;

                                PatternRewriter &rewriter) const override {
    // Removes the custom call with sharding op if the operand type is the
    // same as the result type.
    if (op->hasAttr(kShardingAttr) && op.getCallTargetName() == kShardingName &&
        op.getNumOperands() == 1 && op.getNumResults() == 1 &&
        op.getOperands().front().getType() ==
            op.getResults().front().getType()) {

                                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,