constant_op.getResult().hasOneUse())
      continue;
    // Do not duplicate constant op if the size is too large.
    // 32 is chosen to be larger than all constants useful for shape references,
    // while not too large to possibly significantly increase model size.
    if (constant_op.getValue().getNumElements() > 32) continue;
    while (!constant_op.getResult().hasOneUse()) {

};

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

          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(

          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

  // variable. These `tf.TPUPartitionedInputV2` will be removed when rewriting
  // the operands.
  for (Value result : parallel_execute.getExecuteOutputs()) {
    if (!result.hasOneUse()) continue;
    auto assign_var =
        llvm::dyn_cast<TF::AssignVariableOp>(*result.getUsers().begin());
    if (!assign_var || assign_var.getValue() != result) continue;
    auto partitioned_input =

    execute_launches.reserve(compile_launch.getNumResults() - 1);
    for (Value program_result :
         llvm::drop_begin(compile_launch.getResults(), 1)) {
      if (!program_result.hasOneUse()) return;
      Operation* user = *program_result.user_begin();
      auto execute = llvm::dyn_cast<TF::TPUExecuteOp>(user);
      if (!execute) return;
      auto execute_launch =

    if (!pre_quantized.hasOneUse()) {
      return failure();
    }

    op.emitWarning("Remove trivial `rescale` op. Please fix the source graph.");

    llvm::SmallVector<Type, 4> new_output_types;
    for (auto result : def->getResults()) {
      if (result.hasOneUse() && *result.getUsers().begin() == op) {

    Operation* defining_op = curr_operand->get().getDefiningOp();

    if (llvm::isa_and_nonnull<TF::ConstOp>(defining_op)) {
      // No need to clone if this is the only use.
      if (defining_op->hasOneUse()) {
        LLVM_DEBUG(llvm::dbgs()
                   << "Not duplicating constant operand since it has only one "
                      "usage. Op: "
                   << curr_operand->getOperandNumber()