// cases, we can convert this op to other TF Ops, which in later passes
// properly convert to TF Lite ops.
struct ConvertTFEinsumOp : public OpRewritePattern<TF::EinsumOp> {
 public:
  explicit ConvertTFEinsumOp(MLIRContext* context)
      : OpRewritePattern<TF::EinsumOp>(context) {}

  LogicalResult matchAndRewrite(TF::EinsumOp op,

    // For the uniform opset, it is requested to maintain the BatchMatmul logic.
    // For the TF opset, since we need to test the effect we remain it as a
    // future work.
    patterns.add<TF::ConvertTFEinsumOp>(ctx);
  }

  if (failed(applyPatternsAndFoldGreedily(func, std::move(patterns)))) {
    func.emitError() << "quant-prepare-lifting failed.";
    signalPassFailure();
  }
}

}  // namespace

void TransformEinsumPass::runOnOperation() {
  RewritePatternSet patterns(&getContext());
  auto func = getOperation();

  patterns.add<ConvertTFEinsumOp>(&getContext());
  (void)applyPatternsAndFoldGreedily(func, std::move(patterns));
}

}  // namespace

LogicalResult ConvertTFEinsumOp::matchAndRewrite(
    TF::EinsumOp op, PatternRewriter& rewriter) const {
  if (op->getNumOperands() == 1) {

                       ReorderFakeQuantPattern<TF::TransposeOp>>(ctx);
  if (unfold_batch_matmul_) {
    TF::PopulateUnrollTfBatchMatMul(ctx, phase_2_patterns);
  }
  phase_2_patterns.add<TF::ConvertTFEinsumOp, ConvertTFStridedSlice,
                       ConvertRfftToRfft2d, RemoveIdentity>(ctx);
  phase_2_patterns.add<ConvertTFConv2D, ConvertTFDepthwiseConv2dNative>(
      ctx, allow_bf16_and_f16_type_legalization_);