Is there a JAX package for dealing with manifold(hyperbolic)-related models?

[nasyxx]

Such as the Hyperbolic neural network or Riemannian adam optimizer.

[adam-hartshorne]

Is this the sort of thing you are looking for?

https://github.com/SaitejaUtpala/rieoptax

[jakubMitura14]

Geomstats also have jax backend , if you are into riemanian statistics worth exploring

[nasyxx]

Is this the sort of thing you are looking for?

https://github.com/SaitejaUtpala/rieoptax

Thanks! Almost yes. It's worthy of reference.

[nasyxx]

Geomstats also have jax backend , if you are into riemanian statistics worth exploring

Thanks, but it seems the JAX backend for geomstata is still work in progress.