Operation* module);

#define GEN_PASS_DEF_GROUPBYDIALECTPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

class GroupByDialectPass
    : public impl::GroupByDialectPassBase<GroupByDialectPass> {
 public:
  void runOnOperation() override;

 private:
  void processFunction(mlir::func::FuncOp func, int& counter,
                       llvm::SmallDenseSet<StringRef>& dialects,