What hyperbolix actually is

hyperbolix is an open‑source library for hyperbolic deep learning in JAX, built on Flax NNX, and is described as a “comprehensive, general-purpose hyperbolic deep learning library in JAX” [[1]](http://arxiv.org/abs/2609.28248v1). It’s released under the MIT license [[1]](http://arxiv.org/abs/2609.28248v1).

At a high level, it gives you:

  • A common interface over six manifolds [[1]](http://arxiv.org/abs/2609.28248v1):
  • Euclidean space
  • Poincaré ball
  • Hyperboloid
  • κ‑stereographic model
  • Mixed‑curvature product spaces
  • Proper velocity space
  • Families of layers implemented for these manifolds [[1]](http://arxiv.org/abs/2609.28248v1):
  • Linear layers
  • Convolutions
  • Attention
  • Normalization
  • Positional encoding
  • Regression
  • Vector quantization
  • Building blocks that “span methods ranging from Ganea’s original hyperbolic neural networks to recent fully hyperbolic architectures such as Hypformer and Lorentzian ResNet” [[1]](http://arxiv.org/abs/2609.28248v1).
  • Riemannian optimizers implemented as Optax transformations, wrapped distributions, and hyperbolic dimensionality‑reduction techniques [[1]](http://arxiv.org/abs/2609.28248v1).

The library is written in “idiomatic JAX”: manifolds are stateless, curvature is passed in at call time, and manifold operations act on single points, with separate mechanisms enabling batch operations [[1]](http://arxiv.org/abs/2609.28248v1).

For production‑oriented builders, the main implication is that you can keep a JAX/Flax mental model (pure functions, explicit state, Optax for optimization) while introducing non‑Euclidean geometry via well‑scoped abstractions rather than hand‑rolling manifold math.

Manifolds as first‑class configuration

hyperbolix exposes six manifolds through a common interface [[1]](http://arxiv.org/abs/2609.28248v1). The paper does not spell out the exact methods on that interface, but it does say:

  • Manifolds are stateless.
  • Curvature is passed at call time.
  • Manifold operations are defined on single points, with separate support for batching [[1]](http://arxiv.org/abs/2609.28248v1).

In practice, that means any agentic system or model orchestration you already have in JAX can, in principle, treat the manifold as a configuration parameter rather than something baked into the control flow.

From a control‑flow perspective:

  • Stateless manifolds align with JAX’s pure‑function approach. There’s no hidden mutable curvature or running statistics inside a manifold object.
  • Curvature passed at call time means any function that depends on geometry needs curvature as an explicit input. This encourages threading curvature through your transforms and agent steps as a regular argument rather than relying on global state.
  • Pointwise operations imply that batching is handled outside via JAX constructs (e.g., vmap or similar), which keeps the manifold layer composable with the rest of the JAX ecosystem [[1]](http://arxiv.org/abs/2609.28248v1).

For an agent pipeline that already uses JAX and Flax NNX, this suggests an integration pattern where:

  • The agent’s “model factory” or policy builder chooses one of the six manifolds and curvature values as configuration.
  • The core computation graph (layers, attention blocks, etc.) pulls the manifold and curvature as inputs rather than reading them from globals.
  • Vectorization and parallelism stay in JAX land; manifold code just defines the math for individual points.

The paper does not describe an explicit factory or module pattern, but the stateless, call‑time curvature interface is clearly designed to work with the way JAX expects functions to be written [[1]](http://arxiv.org/abs/2609.28248v1).

Layer families: wiring hyperbolic computations end‑to‑end

hyperbolix implements “layer families that cover linear layers, convolutions, attention, normalization, positional encoding, regression, and vector quantization” [[1]](http://arxiv.org/abs/2609.28248v1).

This matters for builders because it implies:

  • You are not restricted to shallow adaptors on top of Euclidean backbones; the library aims to support “fully hyperbolic architectures” [[1]](http://arxiv.org/abs/2609.28248v1).
  • You can build stacks that resemble standard deep learning architectures, with analogues for many familiar components: linear → attention → normalization → positional encodings → task‑specific head, all in a hyperbolic setting.

The authors explicitly say these building blocks “span methods ranging from Ganea’s original hyperbolic neural networks to recent fully hyperbolic architectures such as Hypformer and Lorentzian ResNet” [[1]](http://arxiv.org/abs/2609.28248v1). They do not spell out exact reference implementations of those models, but the presence of the necessary layer types suggests that constructing such architectures is an intended use.

From an orchestration standpoint for agents:

  • Any agent that currently uses Euclidean transformer‑like models in JAX could, in principle, be mirrored with hyperbolic layers for research, ablation, or specialized modules, subject to what the library’s APIs support.
  • Because components like attention and vector quantization are explicitly mentioned, your control flow around discrete routing, codebooks, or memory lookups could potentially be kept structurally similar when moving to a hyperbolic representation, again depending on the exact APIs.

The paper does not describe how these layers compose with Flax NNX modules in detail, but it does state that hyperbolix is built on Flax NNX [[1]](http://arxiv.org/abs/2609.28248v1), which signals that the library is meant to fit into the standard JAX/Flax module and parameter‑tree conventions.

Riemannian optimizers and Optax integration

On the optimization side, the authors say:

“Additionally, hyperbolix contains Riemannian optimizers implemented as optax transformations, wrapped distributions, and hyperbolic dimensionality-reduction techniques.” [[1]](http://arxiv.org/abs/2609.28248v1)

The key engineering detail here is Optax transformations: Optax is the standard optimizer library in the JAX ecosystem. Implementing Riemannian optimizers as Optax transformations means:

  • You can keep your existing training loops, update functions, and agent training infrastructure that already assumes “optimizer is an Optax transformation”.
  • Switching between Euclidean and Riemannian optimization becomes a matter of swapping which Optax transformation you plug into your training step, assuming the rest of your stack (model parameters, manifolds) is wired correctly.

The paper does not enumerate which specific Riemannian optimization algorithms are provided, nor their performance, but the integration hint is concrete: they are packaged in the same way your current Optax SGD/Adam/etc. might be [[1]](http://arxiv.org/abs/2609.28248v1).

The mention of “wrapped distributions” and “hyperbolic dimensionality-reduction techniques” [[1]](http://arxiv.org/abs/2609.28248v1) indicates that the library goes beyond just forward layers and optimizers to tooling for probabilistic modeling and embedding compression in hyperbolic settings. The paper does not give further detail on their APIs or algorithms.

Numerical reliability and the hyperboloid

The most concrete low‑level engineering claim in the paper is about numerical stability on the hyperboloid manifold:

“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.” [[1]](http://arxiv.org/abs/2609.28248v1)

The library addresses this as follows:

“hyperbolix replaces these subtractions with cancellation-free formulas that stay accurate in float32 at distances where prior implementations return NaN.” [[1]](http://arxiv.org/abs/2609.28248v1)

This has direct implications for production:

  • If your agent relies on hyperboloid‑based distances (e.g., for retrieval, clustering, or routing decisions in embedding space), you need those distances to be well‑behaved under float32, which is standard in accelerators.
  • The authors claim that earlier implementations can return NaNs at large distances from the origin due to catastrophic cancellation, while their “cancellation‑free formulas” avoid that, at least for the tested ranges [[1]](http://arxiv.org/abs/2609.28248v1).

They back this with a testing strategy:

“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.” [[1]](http://arxiv.org/abs/2609.28248v1)

Two important engineering details emerge:

  1. Reference checks: For each “checked operation” (the paper does not enumerate them), there’s either a closed‑form NumPy/SciPy equivalent derived from source papers or a finite‑difference approximation used as oracle [[1]](http://arxiv.org/abs/2609.28248v1).
  1. Precision coverage: Both float32 and float64 are tested [[1]](http://arxiv.org/abs/2609.28248v1).

In an agent system that uses JAX on GPU or TPU, float32 is typically the baseline. Having explicit coverage for float32 and float64 gives you a clearer expectation about where the library has been validated, even though the paper does not give numeric error bounds or benchmarks.

For control flow, this suggests:

  • You can more safely rely on hyperboloid distances and other two‑point operations for decision‑making logic (e.g., which memory to read, which branch to execute) without having to wrap everything in NaN‑guards as aggressively as you might with naive hyperbolic implementations.
  • If your agent needs to run in float64 (e.g., for research or CPU‑bound workflows), the same operations have been tested in that precision as well [[1]](http://arxiv.org/abs/2609.28248v1).

How this fits in an agentic JAX stack

Putting the stated pieces together, a reasonable integration pattern looks like:

  • Geometry layer:
  • Choose a manifold from the six available (Euclidean, Poincaré ball, hyperboloid, κ‑stereographic, mixed‑curvature product space, proper velocity space) [[1]](http://arxiv.org/abs/2609.28248v1).
  • Thread curvature as an explicit configuration parameter through your model and agent modules.
  • Model layer:
  • Build or adapt models using the provided hyperbolic layer families (linear, convolution, attention, normalization, positional encoding, regression, vector quantization) [[1]](http://arxiv.org/abs/2609.28248v1).
  • For architectures conceptually similar to Ganea‑style hyperbolic networks, Hypformer, or Lorentzian ResNet, use these building blocks as the primary components since the paper states they span those methods [[1]](http://arxiv.org/abs/2609.28248v1).
  • Optimization layer:
  • Instantiate Riemannian optimizers through Optax transformations supplied by hyperbolix [[1]](http://arxiv.org/abs/2609.28248v1).
  • Plug them into your existing JAX training loop where an Optax transformation is expected.
  • Inference and decision layer:
  • For hyperboloid‑based computations, rely on hyperbolix’s cancellation‑free distance and other two‑point operations for stable float32 behavior far from the origin [[1]](http://arxiv.org/abs/2609.28248v1).
  • Use these operations in agent routing/selection logic if your system’s control flow depends on geometric relationships.
  • Testing and validation:
  • Leverage the fact that hyperbolix has checked operations against NumPy/SciPy transcriptions or finite differences in float32 and float64 [[1]](http://arxiv.org/abs/2609.28248v1).
  • Align your own regression tests with those operations when they are critical to your agent’s correctness.

All of this sits naturally inside a JAX + Flax NNX + Optax stack, since the library is explicitly built on Flax NNX and its optimizers are implemented as Optax transformations [[1]](http://arxiv.org/abs/2609.28248v1).

Why it matters now for builders

The paper positions hyperbolix as, to the authors’ knowledge, “the first comprehensive, general-purpose hyperbolic deep learning library in JAX” [[1]](http://arxiv.org/abs/2609.28248v1). For builders who:

  • Already standardize on JAX/Flax/Optax for model training and agent components, and
  • Want to experiment with or deploy non‑Euclidean geometry,

hyperbolix offers:

  • A unified, stateless manifold interface with explicit curvature arguments [[1]](http://arxiv.org/abs/2609.28248v1).
  • Layer families that cover a broad slice of building blocks (including attention and vector quantization) instead of just embedding adaptors [[1]](http://arxiv.org/abs/2609.28248v1).
  • Riemannian optimizers in the same abstraction level as your current optimizers (Optax transformations) [[1]](http://arxiv.org/abs/2609.28248v1).
  • Specific numerical work on hyperboloid distances with cancellation‑free formulas that maintain float32 accuracy where earlier implementations fail with NaNs [[1]](http://arxiv.org/abs/2609.28248v1).

This combination is particularly relevant if you care about:

  • Avoiding bespoke, fragile hyperbolic math in critical control‑flow paths.
  • Keeping your agent stack consistent with JAX’s functional style and the Flax NNX module ecosystem.

The library is available under the MIT license at a URL referenced in the paper [[1]](http://arxiv.org/abs/2609.28248v1).

What is not documented

The paper does not document:

  • Any quantitative benchmarks (speed, memory, or accuracy) of hyperbolix versus other implementations.
  • The exact API signatures of manifolds, layers, optimizers, or distributions, beyond high‑level descriptions.
  • Which specific Riemannian optimization algorithms, wrapped distributions, or dimensionality‑reduction methods are implemented.
  • Any downstream application results, datasets, or tasks using hyperbolix.
  • How hyperbolix interacts with distributed training, mixed precision beyond float32/float64, or integration with other JAX ecosystem tools outside Flax NNX and Optax.