};

bool PrepareQuantizePass::SetInputNodesQuantizationParams(func::FuncOp func) {
  StringRef func_name = func.getName();
  auto has_quantize_op = [&](const Value arg) {
    return (arg.hasOneUse() &&
            llvm::isa<quantfork::QuantizeCastOp>(*arg.user_begin()));
  };

  bool need_to_set_input_nodes_quantization_params = false;
  for (const BlockArgument arg : func.getArguments()) {

Value SkipIdentity(Value v, bool allow_other_use,
                   llvm::SmallPtrSet<Operation*, 4>* skipped = nullptr) {
  while (auto result = mlir::dyn_cast<OpResult>(v)) {
    if (!(allow_other_use || v.hasOneUse())) break;
    auto op = result.getDefiningOp();
    if (!llvm::isa<TF::IdentityOp, TF::IdentityNOp>(op)) {
      break;
    }
    v = op->getOperand(result.getResultNumber());
    if (skipped) skipped->insert(op);

        diag << "entities 'mul_value' failed to satisfy constraint: "
                "unsupported dimensions";
      });
    }
    if (!is_dynamic_broadcast &&
        !((*conv_op.getODSResults(0).begin()).hasOneUse())) {
      return rewriter.notifyMatchFailure(mul_op, [&](::mlir::Diagnostic &diag) {
        diag << "entities 'conv' failed to satisfy constraint: has one use";
      });
    }

  LogicalResult matchAndRewrite(BroadcastToOp op,
                                PatternRewriter &rewriter) const override {
    // Only rewrite if the Broadcast has only one consumer.
    if (!op.getOutput().hasOneUse()) return failure();

    Operation *user = *op.getOutput().getUsers().begin();

    auto reshape_op = llvm::dyn_cast_or_null<ReshapeOp>(user);
    if (!reshape_op) return failure();

    auto reshape_type =

          continue;
        }
        auto result_tensor_type = result_type.cast<TensorType>();
        // If the user is the Quantize op, it must be the only user.
        if (result.hasOneUse() &&
            llvm::isa<quantfork::QuantizeCastOp>(*result.user_begin())) {
          auto user =
              llvm::cast<quantfork::QuantizeCastOp>(*result.user_begin());
          outputs_replaced.insert(

    return other_convert_op &&
           converter.isLegal(other_convert_op.getOperand().getType());
  } else if (!converter.isLegal(convert_op.getResult().getType())) {
    if (!convert_op.getResult().hasOneUse()) {
      return false;
    }
    auto other_convert_op = dyn_cast_or_null<OtherConvertOp>(
        *convert_op.getResult().getUsers().begin());
    return other_convert_op &&

  for (int i = 0; i != block.getNumArguments(); ++i) {
    Value arg = block.getArgument(0);
    Type arg_type = arg.getType();
    Value new_arg = arg;
    Location loc = func.getLoc();
    if (arg.hasOneUse() && llvm::isa<QuantizeOp>(*arg.user_begin())) {
      auto quantize_op = llvm::cast<QuantizeOp>(*arg.user_begin());
      auto quantize_output = quantize_op.getOutput();

          continue;
        }
        Type result_ele_type =
            mlir::cast<TensorType>(result.getType()).getElementType();
        // If the user is the QuantizeOp, it must be the only user.
        if (result.hasOneUse() && isa<QuantizeOpT>(*result.user_begin())) {
          auto user = cast<QuantizeOpT>(*result.user_begin());
          outputs_replaced.insert(
              {user.getResult(), enumerated_result.index()});

      return failure();
    }
    auto type = mlir::cast<TensorType>(op->getResult(0).getType());
    if (!type || !type.getElementType().isF32()) {
      return failure();
    }
    return success(
        op->hasOneUse() &&
        IsWeightQuantizableFunction(*op->getUses().begin(), type.getRank()));
  }

  // Checks if the operand is second operand of `tf.XlaCallModule` op for

    // There are none type values.
    return;
  }
  if (!tensor_type.getElementType().isF32()) return;

  // If the result is consumed by a quantize op, it has been quantized.
  if (value.hasOneUse() &&
      llvm::isa<TFL::QuantizeOp>(*value.getUsers().begin()))
    return;

  // Add this result to the list to apply the default value.
  values->push_back(value);
}