registry.insert<mlir::func::FuncDialect, mlir::stablehlo::StablehloDialect,
                  mlir::TFL::TFLDialect, mlir::arith::ArithDialect,
                  mlir::TF::TensorFlowDialect, mlir::chlo::ChloDialect>();

  return failed(
      mlir::MlirOptMain(argc, argv, "ODML Converter Driver\n", registry));

  registry.insert<func::FuncDialect>();
  registry.insert<tensor::TensorDialect>();
  registry.insert<shape::ShapeDialect>();
  registry.insert<TF::TensorFlowDialect>();
  registry.insert<chlo::ChloDialect>();

  return registry;
}

class XlaLegalizeTargetsTest : public testing::Test {
 public:
  XlaLegalizeTargetsTest()
      : context_(GetDefaultDialectRegistry()),

                  mlir::quantfork::QuantizationForkDialect,
                  mlir::stablehlo::StablehloDialect,
                  mlir::tf_executor::TensorFlowExecutorDialect,
                  mlir::vhlo::VhloDialect>();
  mlir::mhlo::registerAllMhloDialects(registry);
  mlir::func::registerAllExtensions(registry);
  return failed(
      mlir::MlirOptMain(argc, argv, "StableHLO quant Pass Driver\n", registry));

 public:
  StablehloVhloTypeConverter() : mlir::vhlo::VhloTypeConverter() {
    addConversion([](mlir::Type type) -> mlir::Type {
      if (type.getDialect().getNamespace() ==
          mlir::vhlo::VhloDialect::getDialectNamespace()) {
        return type;
      }
      return {};
    });
    addConversion([](mlir::stablehlo::TokenType token) -> mlir::Type {
      return mlir::vhlo::TokenV1Type::get(token.getContext());

    return "Legalize TF_XlaCallModule Op to stablehlo";
  }
  void getDependentDialects(DialectRegistry &registry) const override {
    registry.insert<stablehlo::StablehloDialect, vhlo::VhloDialect,
                    quant::QuantizationDialect, shape::ShapeDialect>();
  }

  void runOnOperation() override {
    ModuleOp module_op = getOperation();
    RewritePatternSet patterns(&getContext());

      registry.insert<TFL::TensorFlowLiteDialect>();
      registry.insert<arith::ArithDialect>();
      registry.insert<func::FuncDialect>();
      registry.insert<mlir::vhlo::VhloDialect>();
      registry.insert<mlir::stablehlo::StablehloDialect>();
    });

def MergeFusionWithDequantizePass : Pass<"stablehlo-merge-fusion-with-dequantize", "mlir::ModuleOp"> {
  let summary = "Merge quantized conv/dot_general fusion with subsequent dequantize.";
  let dependentDialects = [
    "chlo::ChloDialect",
    "mlir::stablehlo::StablehloDialect",
  ];
}

def UnwrapXlaCallModuleOpPass : Pass<"stablehlo-unwrap-xla-call-module-op", "ModuleOp"> {

                  mlir::quantfork::QuantizationForkDialect,
                  mlir::TFL::TensorFlowLiteDialect, mlir::TF::TensorFlowDialect,
                  mlir::stablehlo::StablehloDialect, mlir::vhlo::VhloDialect>();
  mlir::func::registerAllExtensions(registry);
  context->appendDialectRegistry(registry);

  context->loadDialect<
      mlir::arith::ArithDialect, mlir::func::FuncDialect,

    tfl_dialect_ = module.getContext()
                       ->getOrLoadDialect<mlir::TFL::TensorFlowLiteDialect>();
    vhlo_dialect_ =
        module.getContext()->getOrLoadDialect<mlir::vhlo::VhloDialect>();
    // Right now the TF executor dialect is still needed to build NodeDef.
    module.getContext()
        ->getOrLoadDialect<mlir::tf_executor::TensorFlowExecutorDialect>();
  }