return HasValidWeightOnlyPtqMethod(method.weight_only_ptq(), rank);
  }

  void rewrite(Operation* op, PatternRewriter& rewriter) const override {
    Operation* quantizable_op = *op->getUsers().begin();
    DenseFPElementsAttr attr;
    matchPattern(op->getResult(0), m_Constant(&attr));

    Method method = GetQuantizationMethodOrDefault(quantizable_op);

  // 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 =
        llvm::dyn_cast_or_null<TF::TPUPartitionedInputV2Op>(
            assign_var.getResource().getDefiningOp());

  CPred<"quant::CastQuantizedTypeAttrFromExpressedType($_builder, $0, $1.getType(), " # i # ")">>;

class UsedBy<string op> : Constraint<
  CPred<"llvm::isa<mlir::TFL::" # op # "Op>(*$0.getUsers().begin())">>;

// When the op is passing-through, the output types of the quantized ops need
// to be updated as well. Since the quantize op manages its own type by the

    if (!op->hasAttr(kQuantTraitAttrName) || op->getNumResults() != 1) {
      return failure();
    }

    Value result = op->getResult(0);
    for (auto user : result.getUsers()) {
      if (dyn_cast_or_null<TF::DumpTensorOp>(user)) return failure();
    }

    const FlatSymbolRefAttr f_attr = GetFuncAttr(op);
    if (!f_attr.getValue().starts_with(kCompositeFuncPrefix)) return failure();

    // Match only when there is one user of the dequantize op.
    if (!op.getResult().hasOneUse()) {
      return failure();
    }

    for (Operation* op_with_region : op.getResult().getUsers()) {
      // Among the ops with regions, only reduce_window op is supported for now.
      if (!isa<mlir::stablehlo::ReduceWindowOp>(op_with_region)) {
        return failure();
      }

  // `dequantize_op`'s users.
  FailureOr<SmallVector<Operation*>> CollectCandidateOps(
      DequantizeOpT dequantize_op) const {
    auto users = dequantize_op->getResult(0).getUsers();
    return SmallVector<Operation*>(users.begin(), users.end());
  }

  // Collects all candidate ops for quantization, which is the operand of

    auto result = execute_output.value();
    if (!result.hasOneUse()) {
      if (VLOG_IS_ON(2)) {
        bool any_user_is_assign = false;
        for (auto result_user : result.getUsers()) {
          any_user_is_assign |= llvm::isa<TF::AssignVariableOp>(result_user);
        }
        if (any_user_is_assign) {
          VLOG(2) << "TPUMergeVariablesWithExecutePass: Skipping output "

    auto ele = op->getResult(res).getType().cast<ShapedType>().getElementType();
    if (ele.isa<FloatType>() && SetResultParams(op, res, params)) {
      auto users = op->getResult(res).getUsers();
      *changed |= !users.empty();
      new_items->append(users.begin(), users.end());
    }
  }
  return success();
}

int QuantizeContext::StatesManager::InitializeState(

def IsRank4Tensor : Type<HasAnyRankOfPred<[4]>, "Rank 4 tensor">;

// Checks if all the users is ReadVariableOp.
def HasOnlyReadVariableOpUsers : Constraint<
  CPred<"llvm::all_of($0.getUsers(), [](mlir::OpOperand op) { "
        "return llvm::isa<mlir::TF::ReadVariableOp>(op.getOwner()); })">>;

def CopyAttrs: NativeCodeCallVoid<"CopyDeviceAndUnderscoredAttributesAdaptor($0, $1)">;

// type if provided. If there is no user or user of type, return nullptr.
template <typename T = Operation*>
Operation* FindUserOfType(Operation* op) {
  for (Operation* user : op->getUsers()) {
    if (isa<T>(user)) {
      return user;
    }
  }
  return nullptr;
}

// Returns the first user of the given operation, optionally of the given