#include "tensorflow/compiler/mlir/lite/transforms/passes.h.inc"

struct ConvertTensorListPopBack
    : public OpRewritePattern<TF::TensorListPopBackOp> {
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(TF::TensorListPopBackOp op,
                                PatternRewriter& rewriter) const override {
    // It is currently not possible to easily pack the output of a multi-result

struct RemoveVolatileOps : public OpRewritePattern<DequantizeOp> {
  explicit RemoveVolatileOps(MLIRContext* context)
      : OpRewritePattern<DequantizeOp>(context, 1) {}

  LogicalResult matchAndRewrite(DequantizeOp op,
                                PatternRewriter& rewriter) const override {
    auto input_op = op.getInput().getDefiningOp();
    if (auto q = llvm::dyn_cast_or_null<QuantizeOp>(input_op)) {

class ConvertResultsBroadcastableShapeOp : public RewritePattern {
 public:
  ConvertResultsBroadcastableShapeOp(MLIRContext* context)
      : RewritePattern(MatchAnyOpTypeTag(), 1, context) {}

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

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

// pass.
struct RemoveIdentity : public OpRewritePattern<TF::IdentityOp> {
  using OpRewritePattern<TF::IdentityOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(TF::IdentityOp identity,
                                PatternRewriter &rewriter) const override;
};

}  // namespace quant
}  // namespace mlir

struct SqueezeReshapesAroundBroadcastOp
    : public OpRewritePattern<TFL::BroadcastToOp> {
  using OpRewritePattern<TFL::BroadcastToOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(TFL::BroadcastToOp tfl_broadcast_to_op,
                                PatternRewriter &rewriter) const override {
    auto loc = tfl_broadcast_to_op->getLoc();

    // Match the
    // Reshape(

#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"  // IWYU pragma: keep
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"

namespace mlir {
namespace odml {

LogicalResult ConvertCustomCallOp::matchAndRewrite(
    mhlo::CustomCallOp mhlo_custom_call, OpAdaptor adaptor,
    ConversionPatternRewriter& rewriter) const {
  auto tfl_custom = rewriter.create<TFL::CustomOp>(

namespace mlir {
namespace TF {

namespace {

struct FuseParallelMapAndBatch : public OpRewritePattern<BatchDatasetV2Op> {
  using OpRewritePattern<BatchDatasetV2Op>::OpRewritePattern;

  LogicalResult matchAndRewrite(BatchDatasetV2Op op,
                                PatternRewriter &rewriter) const override {
    auto batchInputDataset = op.getInputDataset();

struct ConstantFoldQuantizableOperands : public RewritePattern {
 public:
  explicit ConstantFoldQuantizableOperands(MLIRContext* context)
      : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}
  LogicalResult matchAndRewrite(Operation* op,
                                PatternRewriter& rewriter) const override;
};

}  // namespace quant
}  // namespace mlir

// Fold quantized i32 (normally bias) into their float values.
struct FoldQuantizedI32ToFloat : public OpRewritePattern<TFL::DequantizeOp> {
  using OpRewritePattern<TFL::DequantizeOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(TFL::DequantizeOp dequant_op,
                                PatternRewriter& rewriter) const override {
    // We only fold i32 -> float pattern.
    auto input = dequant_op.getInput().getDefiningOp();

}

// Transform pad(transpose(x)) to transpose(pad(x))
struct TransposeCommuteWithPad : public OpRewritePattern<stablehlo::PadOp> {
  using OpRewritePattern<stablehlo::PadOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(stablehlo::PadOp pad_op,
                                PatternRewriter& rewriter) const override {
    Value pad_input = pad_op.getOperand();
    RankedTensorType pad_type = pad_op.getType().cast<RankedTensorType>();