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

  context->loadDialect<

    // This is a brief description of the pass.
    return "Quantize composite functions with QDQ input/outputs.";
  }

  void getDependentDialects(DialectRegistry& registry) const override {
    registry.insert<TF::TensorFlowDialect, quant::QuantizationDialect,
                    quantfork::QuantizationForkDialect>();
  }

 private:
  void runOnOperation() override;

  bool enable_legacy_weight_only_;

  ConversionTarget target(*context);
  target.addLegalDialect<arith::ArithDialect>();
  target.addLegalDialect<func::FuncDialect>();
  target.addLegalDialect<TF::TensorFlowDialect>();
  target.addLegalOp<ModuleOp>();
  target.addLegalOp<func::FuncOp>();
  target.addIllegalOp<TF::XlaConvV2Op>();
  target.addIllegalOp<TF::XlaGatherOp>();

  RewritePatternSet patterns(context);

               .cast<mlir::BoolAttr>()
               .getValue();
  };

  ConversionTarget target(*context);
  target.addDynamicallyLegalDialect<TF::TensorFlowDialect>(is_legal);
  target.addIllegalOp<TF::EmptyTensorListOp, TF::TensorListFromTensorOp,
                      TF::TensorListGetItemOp, TF::TensorListLengthOp,

}

struct EmbeddingPipeliningPass
    : public ::impl::EmbeddingPipeliningPassBase<EmbeddingPipeliningPass> {
  void getDependentDialects(mlir::DialectRegistry& registry) const override {
    registry.insert<TF::TensorFlowDialect>();
  }

  void runOnOperation() override;
};

bool UseEmbeddingPipelining(ModuleOp& module) {
  // Enable automated pipelining pass unless:
  // 1. The user disables it via flag, or