hyperbolix 0.3.0

Hyperbolic Deep Learning in JAX

Hyperbolix

Hyperbolic Deep Learning in JAX

Pure JAX implementation of hyperbolic deep learning with manifold operations, neural network layers, and Riemannian optimizers. Built with Flax NNX and Optax.

Features

• 🌐 3 Manifolds: Euclidean, Poincaré Ball, Hyperboloid
• 🧠 13+ Neural Network Layers: Linear, convolutional (2D/3D), regression
• ⚡ 4 Hyperbolic Activations: ReLU, Leaky ReLU, Tanh, Swish
• 📈 Riemannian Optimizers: RAdam and RSGD with automatic manifold detection
• 🚀 Pure JAX/Flax NNX: vmap-native API, JIT-compatible (10-100x speedup)
• ✅ 1,400+ tests passing with comprehensive benchmark suite

Quick Start

import jax.numpy as jnp
from hyperbolix.manifolds import poincare
from hyperbolix.nn_layers import HypLinearPoincare
from flax import nnx

# Manifold operations
x = jnp.array([0.1, 0.2])
y = jnp.array([0.3, -0.1])
distance = poincare.dist(x, y, c=1.0, version_idx=0)

# Neural network layer
layer = HypLinearPoincare(
    manifold_module=poincare,
    in_dim=128,
    out_dim=64,
    rngs=nnx.Rngs(0)
)
output = layer(x_batch, c=1.0)

Installation

git clone https://github.com/hyperbolix/hyperbolix.git
cd hyperbolix
uv sync  # or: pip install -e .

Requirements: Python 3.12+, JAX 0.4.20+, Flax 0.8.0+, Optax 0.1.7+

Documentation

📖 Full Documentation

• Getting Started - Installation and first examples
• User Guides - Manifolds, layers, optimizers, batching, numerical stability
• API Reference - Complete API documentation
• Developer Guide - Development setup and workflows

Build docs locally: uv run mkdocs serve

Key Concepts

Pure functional design: No stateful classes, curvature passed at call time

import hyperbolix.manifolds.poincare as poincare
dist = poincare.dist(x, y, c=1.0, version_idx=0)  # (dim,) → scalar

vmap-native API: Functions operate on single points, use jax.vmap for batching

# Batch operations
distances = jax.vmap(poincare.dist, in_axes=(0, 0, None, None))(
    x_batch, y_batch, 1.0, 0
)

Citation

@software{hyperbolix2026,
  title = {Hyperbolix: Hyperbolic Deep Learning in JAX},
  author = {Klein, Timo and Lang, Thomas},
  year = {2026},
  url = {https://github.com/hyperbolix/hyperbolix}
}

References

Implements methods from:

• Ganea et al. (2018): Hyperbolic Neural Networks
• Bécigneul & Ganea (2019): Riemannian Adaptive Optimization
• Nagano et al. (2019): Wrapped Normal Distribution on Hyperbolic Space
• Shimizu et al. (2020): Hyperbolic Neural Networks++
• Bdeir et al. (2023): Fully Hyperbolic CNNs
• Bdeir et al. (2025): Robust Hyperbolic Learning

See individual module docstrings for detailed references.

Contributing

Contributions welcome! See DEVELOPER_GUIDE.md for setup and guidelines.

For bugs or questions, open an issue.

License

MIT License. See LICENSE for details.