Hyperbolic Graph Neural Networks in JAX

This repository contains a full implementation of Hyperbolic Graph Attention Networks (HGATs) using the Lorentz (Hyperboloid) Model in JAX. It's designed for modeling hierarchical data structures (like knowledge graphs) and features a Host-to-Device paging strategy for scaling on single-GPU hardware.

• Docs: Read the comprehensive Documentation Directory for deep dives into Riemannian optimization, the Lorentz model, hyperbolic InfoNCE loss, and attention aggregation.
• Walkthrough: See the Implementation Walkthrough for a guide bridging the math with the Python codebase, along with training visualizations showcasing the clustered hierarchy.

Installation

This project utilizes uv for dependency management and requires Python 3.12+.

The main dependencies are:

• jax and jaxlib for hardware-accelerated, differentiable tensor operations.
• matplotlib for stereographic projection visualizations to the Poincaré disk.
• scipy for logging and stable metric computations.
• networkx and pyyaml for structural graph processing and configuration parsing.

System Dependencies

Before syncing the Python environment, ensure you have the following system-level dependencies installed. These are required by certain development packages (like pygraphviz) and are not automatically installed via uv sync:

• Graphviz: Required for graph layout computing and rendering.
  • macOS (Homebrew): brew install graphviz
  • Ubuntu/Debian: sudo apt-get install graphviz libgraphviz-dev

Setup Instructions

1. Ensure you have uv installed. You can install it via the official standalone script:

   curl -LsSf https://astral.sh/uv/install.sh | sh

   macOS Alternative (Homebrew): If you prefer using Homebrew for package management on Mac, you can install it via:

   brew install uv

2. Clone the repository and navigate into it:

   git clone <repository_url>cd hyperbolic

3. Sync the environment and install dependencies using uv:

   uv sync

Running the Demos

The repository includes demonstration scripts to showcase the capabilities of the engine.

End-to-End Demo (demo.py)

This script synthesizes hierarchical data, initializes the network, trains on a single GPU using host-to-device paging, and renders the before/after visualizations.

API Demo (api_demo.py)

This script demonstrates the end-to-end usage of the HyperbolicEngine API on a standard NetworkX graph. It proves that the engine can consume an arbitrary string-labeled graph, automatically extract node relationships (Markov Blankets) without explicit hierarchical tags, and map these nodes onto the Lorentz Manifold. It also shows how to query the continuous space using Riemannian distance to find structurally homologous nodes. To see how nested hierarchical string architectures (like org-charts) can be parsed directly into these NetworkX graphs dynamically, see the YAML Interface Documentation.

To run the demos, navigate to the demo directory:

cd demo
uv run api_demo.py
uv run demo.py

Running api_demo.py will save its visualization to demo/api_poincare_viz.png. Running demo.py will output the training logs and save the final clustered embedding representation to demo/poincare_viz_final.png.

Running Tests

To ensure the core geometric functions maintain mathematical correctness and testing against rank-collapse bugs, you can run the provided Pytest suite:

uv run pytest

License

This project is licensed under the MIT License - see the LICENSE file for details.

Author

• Robert Jones - rljonesiii

Citation

If you use this library in your research, please cite:

@software{jones2026hyperbolic,
author = {Robert Jones},
title = {Hyperbolic Graph Neural Networks in JAX},
url = {https://github.com/rljonesiii/hyperbolic},
year = {2026},
}