};

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

namespace mlir {
namespace TFR {
namespace {
class ConvertConstToTensorConst : public OpRewritePattern<ConstantTensorOp> {
  using OpRewritePattern<ConstantTensorOp>::OpRewritePattern;

 public:
  LogicalResult matchAndRewrite(ConstantTensorOp cst_tensor_op,
                                PatternRewriter &rewriter) const override {
    Location loc = cst_tensor_op.getLoc();
    Type out_type = cst_tensor_op.getType();

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

namespace mlir {
namespace odml {

class ConvertCustomCallOp : public OpConversionPattern<mhlo::CustomCallOp> {
 public:
  using OpConversionPattern::OpConversionPattern;

  LogicalResult matchAndRewrite(
      mhlo::CustomCallOp mhlo_custom_call, OpAdaptor adaptor,
      ConversionPatternRewriter& rewriter) const final;
};

// Ops that have a call_target_name starting with the prefix "custom_call." and

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