Multiplying by edge_weights when doing keras nonparametric UMAP

Within the ParametricUMAP class one particular parameter of interest is parametric_embedding. The documentation described that when set to false, a non-parametric embedding is learned, using the same code as the parametric embedding, which can serve as a direct comparison between parametric and non-parametric embedding using the same optimizer; however, upon close inspection it seems that in the nonparametric case the compute_loss function is slightly different. Namely, the loss is additionally multiplied by the edge_weights:

if not parametric_embedding:
        # multiply loss by weights for nonparametric
        weights_tiled = np.tile(edge_weights, negative_sample_rate + 1)

...

if not parametric_embedding:
            ce_loss = ce_loss * weights_tiled

Do you think someone could provide an intuition for this difference and if it affects our ability to directly compare between parametric and non-parametric embeddings. Specifically, I am hoping to use the ce_loss as a proxy for comparing the performance of PUMAP vs. UMAP on my data.

Start in the ParametricUMAP class at compute_loss and inspect the two non-parametric branches shown in the issue. Compare how edge_weights affect the loss for parametric and non-parametric embeddings, then determine whether the implementation or documentation needs clarification so the intended comparison is well defined.