Positional Encoding

Position-aware layers for hyperbolic Transformers on the hyperboloid. HOPE (hope / HyperbolicRoPE) is a deterministic rotary encoding with no learnable parameters that preserves relative position (⟨HOPE(q,i), HOPE(k,j)⟩_L depends only on i-j); HypformerPositionalEncoding is learnable (HTCLinear + Lorentzian residual). For when to prefer each, see the NN Layers guide.

The lorentz_residual / lorentz_scale functions that back the residual skip connection are documented under Primitives; the LorentzResidual NNX module wrapper is here.

HOPE (Hyperbolic Rotary Positional Encoding)

hyperbolix.nn_layers.hope

hope(
    z: Float[Array, "... seq d_plus_1"],
    positions: Float[Array, seq],
    c: float = 1.0,
    base: float = 10000.0,
    eps: float = 1e-07,
) -> Float[Array, "... seq d_plus_1"]

Hyperbolic Rotary Positional Encoding (HOPE).

Applies RoPE-style rotation to the spatial components of hyperboloid points, then reconstructs the time component to satisfy the manifold constraint. Equivalent to hrc(z, R_{i,Theta}, c, c) where R is a block-diagonal rotation matrix.

Since rotation preserves norms, the Minkowski inner product between encoded points depends only on the relative position offset, giving the standard RoPE relative-position property on the hyperboloid.

Parameters:

Name	Type	Description	Default
z	(Array, shape(..., seq_len, d + 1))	Points on hyperboloid (d must be even).	required
positions	(Array, shape(seq_len))	Integer position indices.	required
c	float	Curvature parameter (default: 1.0).	1.0
base	float	Frequency base for rotation angles (default: 10000.0).	10000.0
eps	float	Numerical stability floor (default: 1e-7).	1e-07

Returns:

Type	Description
(Array, shape(..., seq_len, d + 1))	Rotated points on hyperboloid with curvature c.

References

Chen et al., "Hyperbolic Embeddings for Learning on Manifolds" (HELM), 2024.

Source code in hyperbolix/nn_layers/hyperboloid_positional.py

def hope(
    z: Float[Array, "... seq d_plus_1"],
    positions: Float[Array, "seq"],
    c: float = 1.0,
    base: float = 10000.0,
    eps: float = 1e-7,
) -> Float[Array, "... seq d_plus_1"]:
    """Hyperbolic Rotary Positional Encoding (HOPE).

    Applies RoPE-style rotation to the spatial components of hyperboloid
    points, then reconstructs the time component to satisfy the manifold
    constraint. Equivalent to ``hrc(z, R_{i,Theta}, c, c)`` where R is a
    block-diagonal rotation matrix.

    Since rotation preserves norms, the Minkowski inner product between
    encoded points depends only on the *relative* position offset, giving
    the standard RoPE relative-position property on the hyperboloid.

    Parameters
    ----------
    z : Array, shape (..., seq_len, d+1)
        Points on hyperboloid (d must be even).
    positions : Array, shape (seq_len,)
        Integer position indices.
    c : float, optional
        Curvature parameter (default: 1.0).
    base : float, optional
        Frequency base for rotation angles (default: 10000.0).
    eps : float, optional
        Numerical stability floor (default: 1e-7).

    Returns
    -------
    Array, shape (..., seq_len, d+1)
        Rotated points on hyperboloid with curvature c.

    References
    ----------
    Chen et al., "Hyperbolic Embeddings for Learning on Manifolds" (HELM), 2024.
    """
    spatial_SD = z[..., 1:]  # (..., S, D) where S=seq, D=spatial dim
    d = spatial_SD.shape[-1]
    dtype = spatial_SD.dtype

    # Frequency schedule: theta_i = 1 / base^(2i/d). Build the index grid in the
    # input dtype: bare jnp.arange is int64 under global jax_enable_x64, and the
    # subsequent division would promote the whole encoding to float64.
    freqs_F = 1.0 / (base ** (jnp.arange(0, d, 2, dtype=dtype) / d))  # (F,) where F = d//2
    angles_SF = positions[:, None].astype(dtype) * freqs_F[None, :]  # (S, F)
    cos_SF = jnp.cos(angles_SF)
    sin_SF = jnp.sin(angles_SF)

    # Rotate spatial components (interleaved pairs)
    rotated_SD = _apply_rotary_interleaved(spatial_SD, cos_SF, sin_SF)  # (..., S, D)

    # Reconstruct time: t = sqrt(||rotated||^2 + 1/c)
    norm_sq_S1 = jnp.sum(rotated_SD**2, axis=-1, keepdims=True)  # (..., S, 1)
    time_S1 = jnp.sqrt(jnp.maximum(norm_sq_S1 + 1.0 / c, eps))  # (..., S, 1)

    return jnp.concatenate([time_S1, rotated_SD], axis=-1)  # (..., S, A)

hyperbolix.nn_layers.HyperbolicRoPE

HyperbolicRoPE(
    dim: int,
    max_seq_len: int = 2048,
    base: float = 10000.0,
    eps: float = 1e-07,
)

Bases: Module

NNX module wrapper for HOPE (Hyperbolic Rotary Positional Encoding).

This is a stateless module (no learnable parameters) that wraps the functional :func:hope for convenient use in NNX model definitions.

Parameters:

Name	Type	Description	Default
dim	int	Spatial dimension d (must be even).	required
max_seq_len	int	Maximum sequence length (for documentation; not enforced, default: 2048).	2048
base	float	Frequency base for rotation angles (default: 10000.0).	10000.0
eps	float	Numerical stability floor (default: 1e-7).	1e-07

References

Chen et al., "Hyperbolic Embeddings for Learning on Manifolds" (HELM), 2024.

Source code in hyperbolix/nn_layers/hyperboloid_positional.py

def __init__(
    self,
    dim: int,
    max_seq_len: int = 2048,
    base: float = 10000.0,
    eps: float = 1e-7,
):
    self.dim = dim
    self.max_seq_len = max_seq_len
    self.base = base
    self.eps = eps

__call__

__call__(
    z: Float[Array, "... seq d_plus_1"],
    positions: Float[Array, seq],
    c: float = 1.0,
) -> Float[Array, "... seq d_plus_1"]

Apply HOPE positional encoding.

Parameters:

Name	Type	Description	Default
z	(Array, shape(..., seq_len, d + 1))	Points on hyperboloid (spatial dim must equal self.dim, must be even).	required
positions	(Array, shape(seq_len))	Integer position indices.	required
c	float	Curvature parameter (default: 1.0).	1.0

Returns:

Type	Description
(Array, shape(..., seq_len, d + 1))	Rotated points on hyperboloid.

Source code in hyperbolix/nn_layers/hyperboloid_positional.py

def __call__(
    self,
    z: Float[Array, "... seq d_plus_1"],
    positions: Float[Array, "seq"],
    c: float = 1.0,
) -> Float[Array, "... seq d_plus_1"]:
    """Apply HOPE positional encoding.

    Parameters
    ----------
    z : Array, shape (..., seq_len, d+1)
        Points on hyperboloid (spatial dim must equal self.dim, must be even).
    positions : Array, shape (seq_len,)
        Integer position indices.
    c : float, optional
        Curvature parameter (default: 1.0).

    Returns
    -------
    Array, shape (..., seq_len, d+1)
        Rotated points on hyperboloid.
    """
    return hope(z, positions, c, self.base, self.eps)

Hypformer Positional Encoding

hyperbolix.nn_layers.HypformerPositionalEncoding

HypformerPositionalEncoding(
    in_features: int,
    out_features: int,
    *,
    rngs: Rngs,
    epsilon: float = 1.0,
    init_bound: float = 0.02,
    eps: float = 1e-07,
)

Bases: Module

Relative positional encoding from Hypformer.

Computes a position vector via HTCLinear, then combines it with the input using a Lorentzian residual connection:

p = HTCLinear(x)
result = lorentz_residual(x, p, w_y=epsilon, c=c)

where epsilon is a FIXED scalar weight on the position contribution. The Hypformer reference keeps it a plain (non-learnable) tensor fixed at 1.0; making it a trainable parameter is unsafe because gradient descent can drive it below -1, where x + epsilon * p leaves the upper hyperboloid sheet and the abs() in the residual normalizer silently masks the violation (see :func:~hyperbolix.nn_layers.hyperboloid_core.lorentz_residual).

Parameters:

Name	Type	Description	Default
in_features	int	Input ambient dimension (d+1, including time component).	required
out_features	int	Output spatial dimension (d). The HTCLinear output will have ambient dimension d+1 (= out_features + 1), matching the input.	required
rngs	Rngs	Random number generators for parameter initialization.	required
epsilon	float	Fixed scalar weight for the position encoding contribution (default: 1.0, matching the Hypformer reference). Must be >= 0 so the Lorentzian residual stays on the upper hyperboloid sheet.	1.0
init_bound	float	Bound for HTCLinear uniform weight initialization (default: 0.02).	0.02
eps	float	Numerical stability floor for lorentz_residual (default: 1e-7).	1e-07

Attributes:

Name	Type	Description
htc_linear	HTCLinear	Linear transformation producing the position encoding vector.
epsilon	float	Fixed scalar weight for the position encoding contribution.
eps	float	Numerical stability parameter.

References

Chen et al., "Hyperbolic Embeddings for Learning on Manifolds" (HELM), 2024. Hypformer paper (citation to be added).

Source code in hyperbolix/nn_layers/hyperboloid_positional.py

def __init__(
    self,
    in_features: int,
    out_features: int,
    *,
    rngs: nnx.Rngs,
    epsilon: float = 1.0,
    init_bound: float = 0.02,
    eps: float = 1e-7,
):
    if epsilon < 0:
        raise ValueError(
            f"epsilon must be >= 0 (got {epsilon}): a negative weight can push the "
            "Lorentzian residual off the upper hyperboloid sheet."
        )
    self.htc_linear = HTCLinear(in_features, out_features, rngs=rngs, init_bound=init_bound)
    self.epsilon = epsilon
    self.eps = eps

__call__

__call__(
    x: Float[Array, "... dim_plus_1"], c: float = 1.0
) -> Float[Array, "... dim_plus_1"]

Apply learnable positional encoding.

Parameters:

Name	Type	Description	Default
x	(Array, shape(..., d + 1))	Points on hyperboloid with curvature c.	required
c	float	Curvature parameter (default: 1.0).	1.0

Returns:

Type	Description
(Array, shape(..., d + 1))	Positionally-encoded points on hyperboloid with curvature c.

Source code in hyperbolix/nn_layers/hyperboloid_positional.py

def __call__(
    self,
    x: Float[Array, "... dim_plus_1"],
    c: float = 1.0,
) -> Float[Array, "... dim_plus_1"]:
    """Apply learnable positional encoding.

    Parameters
    ----------
    x : Array, shape (..., d+1)
        Points on hyperboloid with curvature c.
    c : float, optional
        Curvature parameter (default: 1.0).

    Returns
    -------
    Array, shape (..., d+1)
        Positionally-encoded points on hyperboloid with curvature c.
    """
    p = self.htc_linear(x, c_in=c, c_out=c)  # (..., d+1)
    return lorentz_residual(x, p, w_y=self.epsilon, c=c, eps=self.eps)

Why epsilon is not learnable

The Hypformer reference keeps epsilon a plain (non-trainable) tensor fixed at 1.0. Making it a parameter is unsafe: gradient descent can drive it below -1, where x + epsilon·p leaves the upper hyperboloid sheet and the abs() in the Lorentzian residual normalizer silently masks the violation instead of raising. For the same reason, lorentz_residual's w_y must always be non-negative.

Lorentzian residual connection

hyperbolix.nn_layers.LorentzResidual

LorentzResidual(
    *,
    init_w_y: float = 1.0,
    learnable_weight: bool = True,
    scale: bool = False,
    init_gamma: float = 2.0,
    learnable_scale: bool = False,
    param_dtype: DTypeLike = jnp.float32,
    eps: float = 1e-07,
)

Bases: Module

LResNet residual (skip) connection with a safely-constrained weight.

Computes the weighted Lorentzian midpoint of a skip branch x and a residual branch y (the latter weighted by w_y), optionally rescaling the result along the geodesic ray from the origin (LResNet Eq. 10)::

out = lorentz_residual(x, y, w_y, c)            # weighted Lorentzian midpoint
if scale:
    out = lorentz_scale(out, gamma, c)          # optional Eq. 10 norm control

Both x and y are points on the hyperboloid with curvature c; the caller produces y (e.g. via a conv / linear layer). This mirrors how :class:~hyperbolix.nn_layers.HypformerPositionalEncoding composes a transform with lorentz_residual -- except that layer keeps its weight fixed (a trainable position weight is unsafe there), whereas this module makes a trainable w_y safe via the softplus reparameterization.

Parameters:

Name	Type	Description	Default
init_w_y	float	Initial value of the residual-branch weight w_y (default: 1.0). Must be >= 0 (> 0 when learnable_weight=True, since softplus cannot represent exactly 0).	1.0
learnable_weight	bool	If True (default), w_y is a softplus-constrained nnx.Param; if False it is a fixed Python float.	True
scale	bool	If True, apply the Eq. 10 Klein-geodesic output scaling after the residual (default: False -- the paper frames it as optional, for low-curvature manifolds).	False
init_gamma	float	Initial value of the scaling constant gamma (default: 2.0, the LResNet CIFAR reference value). Must be > 0. Ignored when scale=False.	2.0
learnable_scale	bool	If True, gamma is a softplus-constrained nnx.Param; if False (default) it is a fixed Python float. Ignored when scale=False.	False
param_dtype	DTypeLike	Storage dtype of any learnable raw parameter (default: jnp.float32), pinned so it does not become float64 under global jax_enable_x64.	float32
eps	float	Numerical stability floor passed to lorentz_residual / lorentz_scale (default: 1e-7).	1e-07

Attributes:

Name	Type	Description
w_y_raw	Param or None	Raw (pre-softplus) weight when learnable_weight=True; otherwise None (the fixed value is stored statically).
gamma_raw	Param or None	Raw (pre-softplus) scaling constant when scale and learnable_scale; otherwise None.
use_scale	bool	Whether the Eq. 10 scaling is applied.

References

He, Neil, Menglin Yang, and Rex Ying. "Lorentzian residual neural networks." Proceedings of the 31st ACM SIGKDD Conference on Knowledge Discovery and Data Mining V. 1. 2025. Residual: :func:~hyperbolix.nn_layers.hyperboloid_core.lorentz_residual; Eq. 10 scaling: :func:~hyperbolix.nn_layers.hyperboloid_core.lorentz_scale.

Source code in hyperbolix/nn_layers/hyperboloid_residual.py

def __init__(
    self,
    *,
    init_w_y: float = 1.0,
    learnable_weight: bool = True,
    scale: bool = False,
    init_gamma: float = 2.0,
    learnable_scale: bool = False,
    param_dtype: DTypeLike = jnp.float32,
    eps: float = 1e-7,
):
    if init_w_y < 0:
        raise ValueError(f"LorentzResidual requires init_w_y >= 0, got {init_w_y}")
    if learnable_weight and init_w_y == 0:
        raise ValueError(
            "LorentzResidual with learnable_weight=True requires init_w_y > 0 (softplus cannot represent exactly 0)."
        )
    if init_gamma <= 0:
        raise ValueError(f"LorentzResidual requires init_gamma > 0, got {init_gamma}")

    self.use_scale = scale
    self.eps = eps

    # y-weight: softplus-constrained nnx.Param when learnable, plain float otherwise.
    if learnable_weight:
        self.w_y_raw = nnx.Param(jnp.array(_inv_softplus(init_w_y), dtype=param_dtype))
        self._w_y = None
    else:
        self.w_y_raw = None
        self._w_y = init_w_y

    # Eq. 10 scaling constant gamma (only consulted when use_scale is True).
    if scale and learnable_scale:
        self.gamma_raw = nnx.Param(jnp.array(_inv_softplus(init_gamma), dtype=param_dtype))
        self._gamma = None
    else:
        self.gamma_raw = None
        self._gamma = init_gamma

__call__

__call__(
    x: Float[Array, "... dim_plus_1"],
    y: Float[Array, "... dim_plus_1"],
    c: float = 1.0,
) -> Float[Array, "... dim_plus_1"]

Apply the (optionally scaled) Lorentzian residual connection.

Parameters:

Name	Type	Description	Default
x	(Array, shape(..., d + 1))	Skip / main branch on the hyperboloid with curvature c.	required
y	(Array, shape(..., d + 1))	Residual branch on the hyperboloid with curvature c (weighted by w_y).	required
c	float	Curvature parameter (default: 1.0).	1.0

Returns:

Type	Description
(Array, shape(..., d + 1))	Points on the hyperboloid with curvature c.

Source code in hyperbolix/nn_layers/hyperboloid_residual.py

def __call__(
    self,
    x: Float[Array, "... dim_plus_1"],
    y: Float[Array, "... dim_plus_1"],
    c: float = 1.0,
) -> Float[Array, "... dim_plus_1"]:
    """Apply the (optionally scaled) Lorentzian residual connection.

    Parameters
    ----------
    x : Array, shape (..., d+1)
        Skip / main branch on the hyperboloid with curvature ``c``.
    y : Array, shape (..., d+1)
        Residual branch on the hyperboloid with curvature ``c`` (weighted by
        ``w_y``).
    c : float, optional
        Curvature parameter (default: 1.0).

    Returns
    -------
    Array, shape (..., d+1)
        Points on the hyperboloid with curvature ``c``.
    """
    # Recover w_y > 0 via softplus when learnable; cast to the input dtype so a
    # float32-pinned param does not silently downcast a float64 forward pass.
    w_y = jax.nn.softplus(self.w_y_raw[...]) if self.w_y_raw is not None else self._w_y
    out = lorentz_residual(x, y, w_y=jnp.asarray(w_y, dtype=x.dtype), c=c, eps=self.eps)

    if self.use_scale:
        gamma = jax.nn.softplus(self.gamma_raw[...]) if self.gamma_raw is not None else self._gamma
        out = lorentz_scale(out, gamma=jnp.asarray(gamma, dtype=x.dtype), c=c, eps=self.eps)

    return out

Example

import jax, jax.numpy as jnp
from flax import nnx
from hyperbolix.nn_layers import hope, HyperbolicRoPE
from hyperbolix.manifolds import Hyperboloid

batch, seq_len, d = 4, 16, 8
spatial = jax.random.normal(jax.random.PRNGKey(42), (batch, seq_len, d)) * 0.1
time = jnp.sqrt(jnp.sum(spatial**2, axis=-1, keepdims=True) + 1.0)
z = jnp.concatenate([time, spatial], axis=-1)  # (4, 16, 9)
positions = jnp.arange(seq_len)

z_enc = hope(z, positions, c=1.0)                          # functional interface
rope = HyperbolicRoPE(dim=d, max_seq_len=64, base=10000.0)  # or the NNX module
z_enc = rope(z, positions, c=1.0)
print(z_enc.shape)  # (4, 16, 9)