this->emit_quant_adaptor_ops_ = emit_quant_adaptor_ops;
  }

  void runOnOperation() override;

 private:
  quant::CustomOpMap custom_op_map_;
};

// Cleans up unnecessary QDQ pattern for input/output ops.
class PostQuantizeRemoveQDQPass
    : public impl::PostQuantizeRemoveQDQPassBase<PostQuantizeRemoveQDQPass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PostQuantizeRemoveQDQPass)

  // type for `arg` is `quantized_type`.
  void QuantizeArg(BlockArgument arg, QuantizedType quantized_type);

  // Inserts the Quantize and Dequantize ops (i.e. QDQ) after `value`. The
  // quantized element type for `value` is `quantized_type`.
  void QuantizeValue(Value value, QuantizedType quantized_type, Location loc);

    DenseElementsAttr constant;
    if (auto rhs = bmm_op.getY(); !matchPattern(rhs, m_Constant(&constant))) {
      // The constant may be preceded by QDQs in models with QDQ format, so we
      // should set it to the real constant.
      auto dq = dyn_cast_or_null<DequantizeOp>(rhs.getDefiningOp());
      if (!dq) return failure();

  // specified in this map are subject to quantization.
  CustomOpMap custom_map;

  // If other than kQDQNone, the model is a floating point graph with QDQ ops
  // to be eliminated and fused into quantized kernels.
  QDQConversionMode qdq_conversion_mode = QDQConversionMode::kQDQNone;
};

// Parses the command line flag strings to the CustomOpMap specification.

  let summary = "Restores function name from XlaCallModule op.";
}

def QuantizeCompositeFunctionsPass : Pass<"stablehlo-quantize-composite-functions", "ModuleOp"> {
  let summary = "Quantize composite functions with QDQ input / outputs.";
  let options = [
    Option<"enable_per_channel_quantized_weight_",
        "enable-per-channel-quantized-weight",
        "bool", /*default=*/"true",

// TODO(b/204265523): Removes this pass after the exporting MLIR to SavedModel
// path is available.
std::unique_ptr<OperationPass<ModuleOp>> CreateInsertMainFunctionPass();

// Converts FakeQuant ops to quant.qcast and quant.dcast (QDQ) pairs.
std::unique_ptr<OperationPass<func::FuncOp>> CreateConvertFakeQuantToQdqPass();

// Lifts the quantizable spots as composite functions.
std::unique_ptr<OperationPass<ModuleOp>>
CreateLiftQuantizableSpotsAsFunctionsPass(

        break;
      }
    }
    if (new_out_quant_dim == -1) {
      return transpose_op.emitOpError(
          "new quantization dimension not found in perm");
    }

    // Insert a QDQ pair with the new quantized dimension after TransposeOp
    auto new_qtype = quant::CreateI8F32UniformQuantizedPerAxisType(
        transpose_op.getLoc(), *rewriter.getContext(), aqtype.getScales(),

      // or results. Skip this node for now.
      if (!params) {
        quantized_.erase(op);
        continue;
      }

      // If this is a QDQ conversion only, the op could have a same-scale
      // requirement for the floating point kernel but allow per-axis
      // quantization for the quantized kernel. If the quantized dimension

    // Match Add.
    DenseElementsAttr added_value;
    Value constant_val = add_op.getRhs();
    if (!matchPattern(constant_val, m_Constant(&added_value))) {
      // The constant may be preceded by QDQs in models with QDQ format, so we
      // should set it to the real constant.
      auto dq = dyn_cast_or_null<DequantizeOp>(constant_val.getDefiningOp());
      if (!dq) return failure();

    return "quant-quantize-composite-functions";
  }

  StringRef getDescription() const final {
    // This is a brief description of the pass.
    return "Quantize composite functions with QDQ input/outputs.";
  }

  void getDependentDialects(DialectRegistry& registry) const override {
    registry.insert<TF::TensorFlowDialect, quant::QuantizationDialect,