data_dumper_logger_config.printAfterIfEnabled(
      partitioning_pass.get(), mlir_module_with_add.get(), print_callback);
  EXPECT_EQ(print_callback_count, 3);

  data_dumper_logger_config.printAfterIfEnabled(
      shape_inference_pass.get(), mlir_module_with_add.get(), print_callback);
  EXPECT_EQ(print_callback_count, 4);

  data_dumper_logger_config.printAfterIfEnabled(

      mlir::OpPrintingFlags op_printing_flags = mlir::OpPrintingFlags());

  void printBeforeIfEnabled(mlir::Pass *pass, mlir::Operation *op,
                            PrintCallbackFn print_callback) override;

  void printAfterIfEnabled(mlir::Pass *pass, mlir::Operation *op,
                           PrintCallbackFn print_callback) override;

 private:
  static void DumpMlir(const std::string &filename,

  std::string filename =
      get_filename_(pass_prefix_ + "before_" + pass_name, op);

  if (ShouldPrint(pass, op)) DumpMlir(filename, print_callback);
}

void DataDumperLoggerConfig::printAfterIfEnabled(
    mlir::Pass *pass, mlir::Operation *op, PrintCallbackFn print_callback) {
  std::string pass_name = pass->getName().str();
  std::string filename = get_filename_(pass_prefix_ + "after_" + pass_name, op);

                                              PrintCallbackFn print_callback) {
  if (ShouldPrint(pass, op)) Log(print_callback, pass, op, "before");
}

void BridgeLoggerConfig::printAfterIfEnabled(mlir::Pass* pass,
                                             mlir::Operation* op,
                                             PrintCallbackFn print_callback) {

  // of 'operation' should be dumped *after* the pass 'pass' has been
  // executed. If the IR should be dumped, 'print_callback' should be invoked
  // with the stream to dump into.
  void printAfterIfEnabled(mlir::Pass* pass, mlir::Operation* op,
                           PrintCallbackFn print_callback) override;

  // Returns `true` iff we should log IR for given `pass` and `op`.

  void printBeforeIfEnabled(mlir::Pass* pass, mlir::Operation* op,
                            PrintCallbackFn print_callback) override {
    Dump(pass, print_callback, /*is_before=*/true);
  }

  void printAfterIfEnabled(mlir::Pass* pass, mlir::Operation* op,
                           PrintCallbackFn print_callback) override {
    Dump(pass, print_callback, /*is_before=*/false);
  }

 private:
  int64_t mlir_pass_count_;