};

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

  patterns.add<ConvertTFBatchMatMulToEinsumOp<TF::BatchMatMulOp>,
               ConvertTFBatchMatMulToEinsumOp<TF::BatchMatMulV2Op>>(
      &getContext());
  (void)applyPatternsAndFoldGreedily(func, std::move(patterns));
}

}  // namespace