}

// Print operation with debug info (to get line number info for debugging)
void dump(StringRef message, Operation* op) {
  llvm::dbgs() << message;
  op->print(llvm::dbgs(), OpPrintingFlags().enableDebugInfo(
                              /*enable=*/true, /*prettyForm=*/true));
  llvm::dbgs() << "\n";
}

// Propagates device assignment inside a function.

  }

  // Export StableHLO MLIR as output
  std::string result;
  llvm::raw_string_ostream os(result);
  OpPrintingFlags printing_flags;
  if (debug_info) {
    printing_flags.enableDebugInfo();
  }
  if (elide_large_elements_attrs) {
    printing_flags.elideLargeElementsAttrs();
  }
  module.print(os, printing_flags);
  os.flush();

  output->os() << result;
  output->keep();