module.walk([&](func::FuncOp f) {
    if (f->hasAttr(kFromXlaCallModuleAttrName)) {
      f->erase();
    }
  });
}

class XlaCallModuleSerializationPass
    : public impl::XlaCallModuleSerializationPassBase<
          XlaCallModuleSerializationPass> {
 public:
  void runOnOperation() override {
    mlir::ModuleOp module = getOperation();
    mlir::SymbolTableCollection symbol_table;