namespace {

using QuantizationComponentSpec =
    tensorflow::quantization::QuantizationComponentSpec;

class EmptyPatternRewriter : public mlir::PatternRewriter {
 public:
  explicit EmptyPatternRewriter(const OpBuilder& other_builder)
      : mlir::PatternRewriter(other_builder) {}
  ~EmptyPatternRewriter() override = default;
};

TEST(TfQuantOpTest, applyUniformQuantization) {
  MLIRContext context;

    SourceMgrDiagnosticHandler sourceMgrHandler(source_manager_, &context_);

    OpBuilder op_builder(&op);
    EmptyPatternRewriter pattern_rewriter(op_builder);

    LogicalResult result =
        Tf2XlaRewriter::RewriteOp(&op, pattern_rewriter,
                                  /*device_type=*/"XLA_CPU_JIT");
    if (!result.succeeded()) {
      return tsl::errors::Internal("Failed to rewrite op");
    }

                                PatternRewriter& rewriter) const override;
};

struct SquaredDifference : public OpRewritePattern<TFL::SquaredDifferenceOp> {
  using OpRewritePattern<TFL::SquaredDifferenceOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(TFL::SquaredDifferenceOp squared_diff_op,
                                PatternRewriter& rewriter) const override;
};

LogicalResult FillAttributesForUniformQuantizedDotOp(
    PatternRewriter& rewriter, Operation* op,
    llvm::StringMap<Attribute>& identifier_to_attr,
    tensorflow::quantization::QuantizationMethod::PresetMethod
        quantization_method,
    bool enable_per_channel_quantization);

LogicalResult FillAttributesForUniformQuantizedConvolutionOp(
    PatternRewriter& rewriter, Operation* op,

  LogicalResult matchAndRewrite(Operation* op,
                                PatternRewriter& rewriter) const override;

 private:
  template <typename Op>
  LogicalResult RewriteEqOp(Operation* op, PatternRewriter& rewriter) const;

  LogicalResult RewriteOp(
      Operation* op, PatternRewriter& rewriter,
      const std::function<bool(ArrayRef<int64_t>, ArrayRef<int64_t>,

// logics.
std::optional<TF::PartitionedCallOp> RegisterOperationsInFuncOp(
    StringRef func_name, PatternRewriter& rewriter, QuantizedType quant_type,
    Value input_val, ShapedType result_type,
    std::function<Operation*(PatternRewriter&, Operation*, Value, ShapedType,
                             QuantizedType)>
        quantization_operations_func) {

  return rewriter.getI64ArrayAttr({stride_h, stride_w});
}

Attribute GetLhsDilationValue(PatternRewriter& rewriter,
                              llvm::StringMap<Attribute>& identifier_to_attr) {
  return rewriter.getI64ArrayAttr({1, 1});
}

Attribute GetRhsDilationValue(PatternRewriter& rewriter,
                              llvm::StringMap<Attribute>& identifier_to_attr) {
  ArrayAttr dilations =

    return failure();
  }

 private:
  LogicalResult InlineRegion(Location loc, PatternRewriter &rewriter,
                             Operation *inline_point, Region *region) const;
};

LogicalResult SimplifySCFIfOp::InlineRegion(Location loc,
                                            PatternRewriter &rewriter,
                                            Operation *inline_point,

                                       PatternRewriter& rewriter);

  static std::vector<Value> sliceInput(Value value, int batch_size,
                                       Location loc, PatternRewriter& rewriter);

  LogicalResult matchAndRewrite(BatchMatMulOpType op,
                                PatternRewriter& rewriter) const override;
};

#define GEN_PASS_DEF_UNROLLBATCHMATMULPASS

// a static broadcast.
Value broadcastToFeatureDim(Location loc, RankedTensorType result_type,
                            Value value1d, Value shape_value,
                            int64_t feature_dim, PatternRewriter &rewriter) {
  auto dims_type =
      RankedTensorType::get(/*shape=*/{1}, rewriter.getIntegerType(64));
  auto dims = DenseIntElementsAttr::get(dims_type, {feature_dim});
  if (shape_value) {