For small datasets, the original UMAP package is fine: https://umap-learn.readthedocs.io/en/latest/

For large datasets (as the UMAP algorithm scales in exponential compute), you will need to use the GPU-accelerated UMAP from cuML. https://docs.rapids.ai/api/cuml/stable/api/#umap