hyperbolix: Hyperbolic Deep Learning in JAX
Not on the map. Why?
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 .
Key facts. Cited by 1 paper 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
31 of 52 references.
References · 31
- 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,305
- 2017Poincar'e Embeddings for Learning Hierarchical Representationscited by 684
- 2018Hyperbolic Neural Networkscited by 402
- 2019Hyperbolic Image Embeddingscited by 221
- 2018Riemannian Adaptive Optimization Methodscited by 194
- 2020Geoopt: Riemannian Optimization in PyTorchcited by 114
- 2019Constant Curvature Graph Convolutional Networkscited by 108
- 2021Fully Hyperbolic Neural Networkscited by 108
- 2019Mixed-curvature Variational Autoencoderscited by 86
- 2024Compositional Entailment Learning for Hyperbolic Vision-Language Modelscited by 74
- 2023Hyperbolic Deep Learning in Computer Vision: A Surveycited by 73
- 2018geomstats: a Python Package for Riemannian Geometry in Machine Learningcited by 71
19 more
- 2020Differentiating through the Fr'echet Meancited by 53
- 2023Fully Hyperbolic Convolutional Neural Networks for Computer Visioncited by 49
- 2021HoroPCA: Hyperbolic Dimensionality Reduction via Horospherical Projectionscited by 47
- 2022The Numerical Stability of Hyperbolic Representation Learningcited by 47
- 2024Hypformer: Exploring Efficient Transformer Fully in Hyperbolic Spacecited by 43
- 2021CO-SNE: Dimensionality Reduction and Visualization for Hyperbolic Datacited by 28
- 2025HELM: Hyperbolic Large Language Models via Mixture-of-Curvature Expertscited by 26
- 2023Poincar'e ResNetcited by 25
- 2025Hyperbolic Deep Learning for Foundation Models: A Surveycited by 20
- 2024Lorentzian Residual Neural Networkscited by 17
- 2024HyperVQ: MLR-based Vector Quantization in Hyperbolic Spacecited by 13
- 2025PHyCLIP: ℓ1-Product of Hyperbolic Factors Unifies Hierarchy and Compositionality in Vision-Language Representation Learningcited by 7
- 2023HypLL: The Hyperbolic Learning Librarycited by 6
- 2026Hyperbolic Busemann Neural Networkscited by 5
- 2025GGBall: Graph Generative Model on Poincar'e Ballcited by 3
- 2025Understanding and Improving Hyperbolic Deep Reinforcement Learningcited by 3
- 2026Fast and Geometrically Grounded Lorentz Neural Networkscited by 3
- 2026Intrinsic Lorentz Neural Networkcited by 3
- 2022Rieoptax: Riemannian Optimization in JAXcited by 1
Built on by
Citing papers per year: 2026: 1.
1 paper.
Similar, not cited
Close in wording, no citation either way
- 2026Riemannian--Lorentz Fusion of Vision Transformers and State-Space Models
- 2025Hyperbolic Optimization
- 2025Hyper-SET: Designing Transformers via Hyperspherical Energy Minimization
- 2026PolyNODE: Variable-dimension Neural ODEs on M-polyfolds
- 2023Learning Distributions on Manifolds with Free-Form Flows
- 2025Hyperbolic Binary Neural Network
- 2023Flow Matching on General Geometries
- 2026Robot Learning on Discrete Surfaces: Theory and Applications ↗