Hyperbolic embedding with adjustable curvature

[gabrieldernbach]

The user guide contains a lovely example for fitting hyperbolic embeddings
https://umap-learn.readthedocs.io/en/latest/embedding_space.html#bonus-embedding-in-hyperbolic-space
hyperbolic_mapper = umap.UMAP(output_metric='hyperboloid', random_state=42).fit(digits.data)

As I understand the hyperboloid model, the resulting space has a constant curvature of -1.
The hyperbolic variational autoencoder literature of recent years dominantly tunes the curvature to the data, yielding convincing results.
I was wondering if there is a simple extension to the "hyperboloid" metric that would allow us to set the curvature as well in umap.

umap/umap/distances.py

Line 206 in ebe5051

def hyperboloid_grad(x, y):

Checking git-blame @lmcinnes last touched this, maybe you can forward me to the right person?

As the guide explicitly mentions the numerical issues of other models, there is a recent contribution from the VAE literature, that may be of interest in that regard https://arxiv.org/pdf/2209.15217.pdf

[lmcinnes]

It's not actually obvious to me how to do this; if you have some references for the hyperboloid model of hyperbolic space under different curvatures I could look into that. Otherwise I would potentially accept pull requests.