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,

        %c = "x.f"(%a)
        %b = "y.f"(%a)
      so that the two "x" dialect instructions are next to each other.
  }];
  let constructor = "TF::CreateOrderByDialectPass()";
}

def GroupByDialectPass : Pass<"tf-group-by-dialect", "mlir::ModuleOp"> {
  let summary = "Groups ops into functions that only contain one dialect.";
  let description = [{
      Factors operations into subroutines such that all functions only