Parametric UMAP run time performance on GPU

Nobody has claimed this yet.

Currently, I am testing the Parametric UMAP on the GPU and I was expecting the performance lift compare to CPU performance.

But it does seem like the run time is not improved at all or even worse.

Here is my setup:
Instance Type - g5.16xlarge
tf version: 2.8
training data set size: 817,614 with 107 columns (after one hot it become 362 columns) - with intention to run 10 times large size later. full data set size is 80MM +
`Parametric parameters:
keras_fit_kwargs = {"callbacks": [
tf.keras.callbacks.EarlyStopping(
monitor='loss',
min_delta=10**-2,
patience=10,
verbose=1,
)
]}

embedder = ParametricUMAP(verbose=True, batch_size = 512, ##512
keras_fit_kwargs = keras_fit_kwargs,
n_training_epochs = 10)

with tf.device(gpus[0].name): ## '/device:GPU:0'
umap_features = embedder.fit_transform(df_train_transformed)`

both GPU and CPU never get to close on the finish line but based on estimated epoche time: CPU estimated 7 mins versus GPU estimated 13 mins.

Any tutorial/best practice to speed up running Parametric UMAP would be highly appreciated! (either GPU or CPU) thanks a lot!

The report names ParametricUMAP.fit_transform as the entry point but no repository file or test. First reproduce the CPU/GPU timing with the supplied dataset and TensorFlow setup, then trace that entry point to identify the runtime bottleneck. Done means a specific, reproducible cause and an agreed performance change or documented best practice.

Written by the indexing model from the issue text.