auto stablehlo_op = builder.create(op_state);
    custom_op.replaceAllUsesWith(stablehlo_op);
    custom_op.erase();
  });
}

std::unique_ptr<OperationPass<func::FuncOp>> CreateTflToStablehloPass() {
  return std::make_unique<TflToStablehloPass>();
}

static PassRegistration<TflToStablehloPass> pass;

}  // namespace odml