};

class RemoveTpuOp : public RewritePattern {
 public:
  explicit RemoveTpuOp(MLIRContext* context)
      : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}

 private:
  LogicalResult matchAndRewrite(Operation* op,
                                PatternRewriter& rewriter) const override {
    // Remove `_tpu_replicate` attributes on each operation first.
    if (op->hasAttr(tensorflow::kTPUReplicateAttr)) {

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

  LogicalResult matchAndRewrite(quantfork::DequantizeCastOp op,
                                PatternRewriter& rewriter) const override {
    auto input_op = op.getArg().getDefiningOp();

// The correct 'bias' broadcast will be synthesized manually.
class ConvertBiasAddOp : public OpRewritePattern<TF::BiasAddOp> {
 public:
  using OpRewritePattern::OpRewritePattern;
  LogicalResult matchAndRewrite(TF::BiasAddOp op,
                                PatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
    tensorflow::TensorFormat data_format;

// converted to an ElementsAttr and then those are packed together using
// tf.Pack op.
struct ConvertConst : public OpConversionPattern<TF::ConstOp> {
  using OpConversionPattern::OpConversionPattern;

  LogicalResult matchAndRewrite(
      TF::ConstOp op, OpAdaptor adaptor,
      ConversionPatternRewriter &rewriter) const override {
    // Verify that the tensor proto contains tensor of type variant and scalar

#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;

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>(

// 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();