namespace TF {

// Create a pass that groups ops into functions that only contain one dialect.
std::unique_ptr<Pass> CreateGroupByDialectPass();

// Register this pass in the global registry of MLIR.
void RegisterGroupByDialectPass();

}  // namespace TF
}  // namespace mlir

    processFunction(func, counter, top_level_dialects_, module);
  });
}

std::unique_ptr<Pass> CreateGroupByDialectPass() {
  return std::make_unique<GroupByDialectPass>();
}

void RegisterGroupByDialectPass() { registerPass(CreateGroupByDialectPass); }

}  // namespace TF