class XlaCallModuleOpToCallOp : public OpRewritePattern<TF::XlaCallModuleOp> {
 public:
  explicit XlaCallModuleOpToCallOp(MLIRContext* context)
      : OpRewritePattern<TF::XlaCallModuleOp>(context) {}

  LogicalResult matchAndRewrite(TF::XlaCallModuleOp op,
                                PatternRewriter& rewriter) const override {
    auto module_op = op->getParentOfType<ModuleOp>();
    SymbolTable symbol_table(module_op);

struct ConvertTFEinsumOp : public OpRewritePattern<TF::EinsumOp> {
 public:
  explicit ConvertTFEinsumOp(MLIRContext* context)
      : OpRewritePattern<TF::EinsumOp>(context) {}

  LogicalResult matchAndRewrite(TF::EinsumOp op,
                                PatternRewriter& rewriter) const override;
};

}  // namespace TF
}  // namespace mlir

#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"

namespace mlir {
namespace quant {

LogicalResult RemoveIdentity::matchAndRewrite(TF::IdentityOp identity,
                                              PatternRewriter &rewriter) const {
  for (Operation *user : identity->getUsers()) {
    // Replace the op with the input if output is only used by TF ops.

}

/// Convert the `tf.IfRegion` op to the `scf.if` op.
class ConvertIfRegionOp : public OpRewritePattern<IfRegionOp> {
 public:
  using OpRewritePattern<IfRegionOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(IfRegionOp op,
                                PatternRewriter& rewriter) const override {
    // Creates the `then` or `else` region of the `scf.if` op. Note that

class TFQuantTypePattern : public ConversionPattern {
 public:
  TFQuantTypePattern(MLIRContext *ctx, TypeConverter &converter)
      : ConversionPattern(converter, MatchAnyOpTypeTag(), 1, ctx) {}

  LogicalResult matchAndRewrite(
      Operation *op, ArrayRef<Value> operands,
      ConversionPatternRewriter &rewriter) const override {
    // This pattern only handle non-UQ, non-const ops.

// during import.
template <typename Op>
struct DropAttributes : public OpRewritePattern<Op> {
  using OpRewritePattern<Op>::OpRewritePattern;

  // Drop the "output_shapes" attribute.
  LogicalResult matchAndRewrite(Op op,
                                PatternRewriter &rewriter) const override {
    bool found = !!op->removeAttr("output_shapes");
    return success(found);
  }
};

 public:
  explicit LowerAddNOp(MLIRContext *context)
      : RewritePattern(AddNOp::getOperationName(), 1, context,
                       {AddV2Op::getOperationName()}) {}

  LogicalResult matchAndRewrite(Operation *op,
                                PatternRewriter &rewriter) const override {
    auto addn_op = cast<AddNOp>(op);

    // TODO(hinsu): Support variant with TensorList type. tf.AddV2 doesn't

template <typename SourceOp>
class PrepareLstmOutputScale : public OpRewritePattern<SourceOp> {
 public:
  explicit PrepareLstmOutputScale(MLIRContext* context)
      : OpRewritePattern<SourceOp>(context) {}
  LogicalResult matchAndRewrite(SourceOp op,
                                PatternRewriter& rewriter) const override {
    operator_property::OpVariant lstm_variant;
    operator_property::OperatorProperty lstm_property;

}

class UnfuseBatchNormTrainingPattern
    : public OpRewritePattern<mhlo::BatchNormTrainingOp> {
 public:
  using OpRewritePattern<mhlo::BatchNormTrainingOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(mhlo::BatchNormTrainingOp bn_op,
                                PatternRewriter &rewriter) const override {
    auto inputs = bn_op.getOperand();
    auto input_type = mlir::dyn_cast<RankedTensorType>(inputs.getType());

 public:
  BFloat16TypePattern(TypeConverter& converter, MLIRContext* ctx)
      : ConversionPattern(converter, MatchAnyOpTypeTag(), /*benefit=*/1, ctx) {}

  LogicalResult matchAndRewrite(
      Operation* op, const ArrayRef<Value> operands,
      ConversionPatternRewriter& rewriter) const override {
    if (getTypeConverter()->isLegal(op)) {
      return failure();
    }