'has_bias': [True, False],
          'has_batch_norm': [True, False],
          'target_opset': [quant_opts_pb2.XLA],
      }])
  )
  @test_util.run_in_graph_and_eager_modes
  def test_qat_conv_model(
      self,
      activation_fn: Optional[ops.Operation],
      has_bias: bool,
      has_batch_norm: bool,
      target_opset: quant_opts_pb2.OpSet,
  ):
    class ConvModel(module.Module):

  // Constructor used by manually creating the pass.
  explicit LiftQuantizableSpotsAsFunctionsDRQPass(
      const QuantMethod quantization_method, const OpSet target_opset,
      const int min_num_elements_for_weights) {
    quantization_method_ = quantization_method;
    target_opset_ = target_opset;
    min_num_elements_for_weights_ = min_num_elements_for_weights;
  }

  LiftQuantizableSpotsAsFunctionsDRQPass(

                        OpSet target_opset)
      : quant_specs_(quant_specs) {
    weight_quantization_ = quant_specs.weight_quantization;
    target_opset_ = target_opset;
  }

  QuantizePass(const QuantizePass& other) : quant_specs_(other.quant_specs_) {
    weight_quantization_ = other.weight_quantization_;
    target_opset_ = other.target_opset_;
  }

  StringRef getArgument() const final {

    : public mlir::OpRewritePattern<quantfork::QuantizeCastOp> {
 public:
  explicit ReplaceQuantizePattern(MLIRContext* context, OpSet target_opset)
      : OpRewritePattern<quantfork::QuantizeCastOp>(context),
        target_opset_(target_opset) {}

 private:
  OpSet target_opset_ = OpSet::TF;

  LogicalResult matchAndRewrite(quantfork::QuantizeCastOp q_op,

// Apply graph optimizations such as fusing and constant folding to prepare
// lifting.
std::unique_ptr<OperationPass<func::FuncOp>> CreatePrepareLiftingPass(
    tensorflow::quantization::OpSet target_opset);

// Lifts the dynamic range quantizable spots as composite functions.
std::unique_ptr<OperationPass<ModuleOp>>
CreateLiftQuantizableSpotsAsFunctionsDRQPass(
    tensorflow::quantization::QuantizationMethod::PresetMethod

std::unique_ptr<OperationPass<ModuleOp>> CreateInsertQuantizedFunctionsPass(
    QuantMethod quantization_method, OpSet target_opset) {
  return std::make_unique<InsertQuantizedFunctionsPass>(quantization_method,
                                                        target_opset);
}

}  // namespace quant

    signalPassFailure();
  }
}

}  // namespace

std::unique_ptr<OperationPass<func::FuncOp>> CreatePrepareLiftingPass(
    const OpSet target_opset) {
  return std::make_unique<PrepareLiftingPass>(target_opset);
}

static PassRegistration<PrepareLiftingPass> pass;

}  // namespace quant

    this->target_ops_ = target_ops;
  }

  void runOnOperation() override;
};

void RaiseCustomOpsPass::runOnOperation() {
  auto fn = getOperation();
  OpBuilder builder(fn.getContext());

  absl::flat_hash_set<std::string> target_op_names(target_ops_.begin(),
                                                   target_ops_.end());

    IntegerAttr num_bits;
    BoolAttr narrow_range;
    ::llvm::SmallVector<Operation *, 4> target_ops;

    target_ops.push_back(op);
    ReorderOp old_reorder_op = ::llvm::dyn_cast<ReorderOp>(op);
    if (failed(findTargetOps(old_reorder_op, rewriter, max, min, input,
                             num_bits, narrow_range, target_ops))) {
      return failure();
    }
    shape = old_reorder_op.getODSOperands(1);

std::unique_ptr<OperationPass<func::FuncOp>> CreateRaiseCustomOpsPass();
std::unique_ptr<OperationPass<func::FuncOp>> CreateRaiseCustomOpsPass(
    const std::vector<std::string>& target_ops);

// Creates raise custom ops pass, which legalize custom ops to TFL::CustomOp
std::unique_ptr<OperationPass<func::FuncOp>> CreateLowerCustomOpsPass();