# hyperbolix: Hyperbolic Deep Learning in JAX

arXiv: https://arxiv.org/abs/2609.28248 · HTML: https://mlsanity.com/paper/2609.28248/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 31 of its 52 references that are in the map.
- First posted on arXiv 23 Sep 2026.

## Builds on (most cited first, 30 of 31)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Poincar'e Embeddings for Learning Hierarchical Representations](https://mlsanity.com/paper/1705.08039/) (2017)
- [Hyperbolic Neural Networks](https://mlsanity.com/paper/1805.09112/) (2018)
- [Hyperbolic Image Embeddings](https://mlsanity.com/paper/1904.02239/) (2019)
- [Riemannian Adaptive Optimization Methods](https://mlsanity.com/paper/1810.00760/) (2018)
- [Geoopt: Riemannian Optimization in PyTorch](https://mlsanity.com/paper/2005.02819/) (2020)
- [Constant Curvature Graph Convolutional Networks](https://mlsanity.com/paper/1911.05076/) (2019)
- [Fully Hyperbolic Neural Networks](https://mlsanity.com/paper/2105.14686/) (2021)
- [Mixed-curvature Variational Autoencoders](https://mlsanity.com/paper/1911.08411/) (2019)
- [Compositional Entailment Learning for Hyperbolic Vision-Language Models](https://mlsanity.com/paper/2410.06912/) (2024)
- [Hyperbolic Deep Learning in Computer Vision: A Survey](https://mlsanity.com/paper/2305.06611/) (2023)
- [geomstats: a Python Package for Riemannian Geometry in Machine Learning](https://mlsanity.com/paper/1805.08308/) (2018)
- [Differentiating through the Fr'echet Mean](https://mlsanity.com/paper/2003.00335/) (2020)
- [Fully Hyperbolic Convolutional Neural Networks for Computer Vision](https://mlsanity.com/paper/2303.15919/) (2023)
- [HoroPCA: Hyperbolic Dimensionality Reduction via Horospherical Projections](https://mlsanity.com/paper/2106.03306/) (2021)
- [The Numerical Stability of Hyperbolic Representation Learning](https://mlsanity.com/paper/2211.00181/) (2022)
- [Hypformer: Exploring Efficient Transformer Fully in Hyperbolic Space](https://mlsanity.com/paper/2407.01290/) (2024)
- [CO-SNE: Dimensionality Reduction and Visualization for Hyperbolic Data](https://mlsanity.com/paper/2111.15037/) (2021)
- [HELM: Hyperbolic Large Language Models via Mixture-of-Curvature Experts](https://mlsanity.com/paper/2505.24722/) (2025)
- [Poincar'e ResNet](https://mlsanity.com/paper/2303.14027/) (2023)
- [Hyperbolic Deep Learning for Foundation Models: A Survey](https://mlsanity.com/paper/2507.17787/) (2025)
- [Lorentzian Residual Neural Networks](https://mlsanity.com/paper/2412.14695/) (2024)
- [HyperVQ: MLR-based Vector Quantization in Hyperbolic Space](https://mlsanity.com/paper/2403.13015/) (2024)
- [PHyCLIP: ℓ1-Product of Hyperbolic Factors Unifies Hierarchy and Compositionality in Vision-Language Representation Learning](https://mlsanity.com/paper/2510.08919/) (2025)
- [HypLL: The Hyperbolic Learning Library](https://mlsanity.com/paper/2306.06154/) (2023)
- [Hyperbolic Busemann Neural Networks](https://mlsanity.com/paper/2602.18858/) (2026)
- [GGBall: Graph Generative Model on Poincar'e Ball](https://mlsanity.com/paper/2506.07198/) (2025)
- [Understanding and Improving Hyperbolic Deep Reinforcement Learning](https://mlsanity.com/paper/2512.14202/) (2025)
- [Fast and Geometrically Grounded Lorentz Neural Networks](https://mlsanity.com/paper/2601.21529/) (2026)
- [Intrinsic Lorentz Neural Network](https://mlsanity.com/paper/2602.23981/) (2026)

## Built on by (most cited first, 1 of 1)

- [Outperforming 1000 Euclidean Layers With Hyperbolic Goal-Conditioned RL](https://mlsanity.com/paper/2610.20009/) (2026)

## Abstract (arXiv)

We present hyperbolix, an open-source library for hyperbolic deep learning in JAX, built on Flax NNX. To our knowledge, it is the first comprehensive, general-purpose hyperbolic deep learning library in JAX. It includes six manifolds with a common interface: Euclidean space, the Poincar'e ball, the hyperboloid, the kappa-stereographic model, mixed-curvature product spaces, and the proper velocity space. We implement layer families that cover linear layers, convolutions, attention, normalization, positional encoding, regression, and vector quantization. These building blocks span methods ranging from Ganea's original hyperbolic neural networks to recent fully hyperbolic architectures such as Hypformer and Lorentzian ResNet. Additionally, hyperbolix contains Riemannian optimizers implemented as optax transformations, wrapped distributions, and hyperbolic dimensionality-reduction techniques. Its API uses idiomatic JAX: Manifolds are stateless, with curvature being passed at call time, while manifold operations act on single points, with vmap enabling batch operations. The precision of every checked operation is tested against a closed-form NumPy/SciPy transcription from the source paper or a finite difference, for both float32 and float64. On the hyperboloid, standard formulas for two-point operations, such as the distance, lose precision far from the origin, because they subtract two large, nearly equal terms. hyperbolix replaces these subtractions with cancellation-free formulas that stay accurate in float32 at distances where prior implementations return NaN. hyperbolix is available under the MIT license at https://github.com/timoklein/hyperbolix .