}

// Returns the function attribute for the given call op which is lifted for
// quantization.
inline FlatSymbolRefAttr GetFuncAttr(TF::PartitionedCallOp call_op) {
  return mlir::dyn_cast<FlatSymbolRefAttr>(call_op.getFAttr());
}

inline FlatSymbolRefAttr GetFuncAttr(TF::XlaCallModuleOp call_op) {
  return call_op->getAttrOfType<FlatSymbolRefAttr>(
      TF::kStablehloEntryFunctionAttrName);
}

  auto new_cluster = builder->create<mlir::tf_device::ClusterOp>(
      cluster.getLoc(), new_cluster_result_types,
      /*operands=*/llvm::ArrayRef<Value>{}, cluster->getAttrs());
  new_cluster.getBody().takeBody(cluster.getBody());

  auto operand_not_in_cluster = [&](mlir::OpOperand& operand) {
    return !new_cluster.getOperation()->isProperAncestor(operand.getOwner());
  };

        exclude = set(exclude)
    if isinstance(obj, BaseModel):
        # TODO: remove when deprecating Pydantic v1
        encoders: Dict[Any, Any] = {}
        if not PYDANTIC_V2:
            encoders = getattr(obj.__config__, "json_encoders", {})  # type: ignore[attr-defined]
            if custom_encoder:
                encoders.update(custom_encoder)
        obj_dict = _model_dump(
            obj,
            mode="json",

    if (cloned == func) continue;
    // Patch up the op attribute to point to the new function.
    for (NamedAttribute attr : op->getAttrs()) {
      auto symref = mlir::dyn_cast<FlatSymbolRefAttr>(attr.getValue());
      if (!symref) continue;
      if (symref.getValue() != func.getName()) continue;
      op->setAttr(attr.getName(),
                  FlatSymbolRefAttr::get(op->getContext(), cloned.getName()));
      break;
    }
  }

      OperationState new_state(candidate_op->getLoc(),
                               candidate_op->getName().getStringRef(), inputs,
                               output_types, candidate_op->getAttrs());
      for (int i = 0; i < candidate_op->getNumRegions(); ++i) {
        new_state.addRegion();
      }
      Operation* quantized_op = rewriter.create(new_state);
      if (candidate_op->getNumRegions() != 0) {

    ModuleOp module_op = op->getParentOfType<ModuleOp>();

    // Ignore ops without quantization method.
    // Consider adding checks for individual methods.
    if (!op->getAttr(kQuantizationMethodAttr)) return failure();

    // Ignore unquantized ops.
    if (!IsQuantizedXlaCallModuleOp(op)) return failure();

    // For weight-only quantization, op should be hybrid quantized.

 private:
  LogicalResult matchAndRewrite(TF::PartitionedCallOp call_op,
                                PatternRewriter& rewriter) const override {
    StringRef function_name =
        mlir::cast<FlatSymbolRefAttr>(call_op.getFAttr()).getValue();
    if (!function_name.starts_with("composite_") ||
        !call_op->hasAttr(kQuantTraitAttrName)) {
      return failure();
    }

    absl::Status check_status;

      new_result_types.push_back(while_op.getResultTypes()[i]);
    }
  }

  auto new_while_op = OpBuilder(while_op).create<WhileOp>(
      while_op.getLoc(), new_result_types, new_operands, while_op->getAttrs());
  new_while_op.getCond().takeBody(while_op.getCond());
  new_while_op.getBody().takeBody(while_op.getBody());

  for (auto i = 0; i < n; ++i) {
    if (!while_op.getResult(i).use_empty()) {

  new_call_op->setAttr(kEntryFuncAttrName,
                       rewriter.getStringAttr(new_ref_func_name.getValue()));
  new_call_op->setAttrs(call_op->getAttrs());
  new_call_op->removeAttr(rewriter.getStringAttr(kQuantTraitAttrName));

  FlatSymbolRefAttr new_func_name_attr =
      FlatSymbolRefAttr::get(rewriter.getContext(), new_ref_func_name);
  new_call_op->setAttr(kEntryFuncAttrName, new_func_name_attr);

                                PatternRewriter& rewriter) const override {
    auto input_op = op.getInput().getDefiningOp();
    if (auto q = llvm::dyn_cast_or_null<QuantizeOp>(input_op)) {
      if (!q->getAttr(mlir::quant::kVolatileOpAttrName)) return failure();

      if (remove_volatile_ops_type == kPreserveInputsAndOutputs) {
        // Don't remove leading and trailing QDQ for PTQ workflow, so the io