Primitives & Helpers

Cross-cutting functions that back the layers on the other pages — the Hyperboloid transformation/regularization components, point-assembly and concatenation helpers, midpoints, residuals, and the statistical reductions used by normalization. They are collected here (rather than scattered across consumer pages) so a given helper has one documented home.

HTC / HRC components

The Hyperbolic Transformation Component (HTC) applies a Euclidean function to the full point; the Hyperbolic Regularization Component (HRC) applies it to the space components only, then reconstructs time. Both support curvature changes (c_in → c_out) and underpin HTCLinear, the HRC* norms, and LorentzConv2D.

hyperbolix.nn_layers.htc

htc(
    x: Float[Array, "... in_dim_plus_1"],
    f_t: Callable[[Float[Array, ...]], Float[Array, ...]],
    c_in: float,
    c_out: float,
    eps: float = 1e-07,
) -> Float[Array, "... out_dim_plus_1"]

Hyperbolic Transformation Component.

Applies a Euclidean linear transformation f_t to the full hyperboloid point (including time component), then maps the result to the hyperboloid with curvature c_out.

Mathematical formula: space = sqrt(c_in/c_out) * f_t(x) time = sqrt(||space||^2 + 1/c_out) output = [time, space]

where f_t takes the full (dim+1)-dimensional input and produces the output spatial components.

Parameters:

Name	Type	Description	Default
x	Array of shape (..., in_dim+1)	Input point(s) on the hyperboloid manifold with curvature c_in. All components (time and spatial) are passed to f_t.	required
f_t	Callable	Euclidean linear transformation applied to the full input. Takes (in_dim+1)-dimensional input and produces out_dim-dimensional output (which becomes the spatial components of the output).	required
c_in	float	Input curvature parameter (must be positive, c > 0).	required
c_out	float	Output curvature parameter (must be positive, c > 0).	required
eps	float	Small value for numerical stability (default: 1e-7).	1e-07

Returns:

Name	Type	Description
y	Array of shape (..., out_dim+1)	Output point(s) on the hyperboloid manifold with curvature c_out.

Notes

• Unlike HRC, f_t operates on the full point including the time component
• f_t's output dimension determines the output spatial dimension
• This is typically used for learnable linear transformations
• The spatial scaling factor sqrt(c_in/c_out) ensures proper curvature transformation

See Also

hrc : Hyperbolic Regularization Component for spatial-only operations. HTCLinear : Module wrapper for htc with learnable linear transformation.

References

Hypformer paper (citation to be added)

Examples:

>>> import jax
>>> import jax.numpy as jnp
>>> from hyperbolix.nn_layers.hyperboloid_core import htc
>>> from hyperbolix.manifolds import Hyperboloid
>>>
>>> # Create a point on the hyperboloid
>>> manifold = Hyperboloid()
>>> x = jnp.array([1.05, 0.1, -0.2, 0.15])
>>> x = manifold.proj(x, c=1.0)
>>>
>>> # Define a linear transformation
>>> W = jax.random.normal(jax.random.PRNGKey(0), (3, 4))
>>> def linear(z):
...     return z @ W.T
>>>
>>> # Apply HTC
>>> y = htc(x, linear, c_in=1.0, c_out=2.0)
>>> y.shape
(4,)  # (3 spatial + 1 time)

Source code in hyperbolix/nn_layers/hyperboloid_core.py

def htc(
    x: Float[Array, "... in_dim_plus_1"],
    f_t: Callable[[Float[Array, "..."]], Float[Array, "..."]],
    c_in: float,
    c_out: float,
    eps: float = 1e-7,
) -> Float[Array, "... out_dim_plus_1"]:
    """Hyperbolic Transformation Component.

    Applies a Euclidean linear transformation f_t to the full hyperboloid point
    (including time component), then maps the result to the hyperboloid with
    curvature c_out.

    Mathematical formula:
        space = sqrt(c_in/c_out) * f_t(x)
        time  = sqrt(||space||^2 + 1/c_out)
        output = [time, space]

    where f_t takes the full (dim+1)-dimensional input and produces the output
    spatial components.

    Parameters
    ----------
    x : Array of shape (..., in_dim+1)
        Input point(s) on the hyperboloid manifold with curvature c_in.
        All components (time and spatial) are passed to f_t.
    f_t : Callable
        Euclidean linear transformation applied to the full input. Takes
        (in_dim+1)-dimensional input and produces out_dim-dimensional output
        (which becomes the spatial components of the output).
    c_in : float
        Input curvature parameter (must be positive, c > 0).
    c_out : float
        Output curvature parameter (must be positive, c > 0).
    eps : float, optional
        Small value for numerical stability (default: 1e-7).

    Returns
    -------
    y : Array of shape (..., out_dim+1)
        Output point(s) on the hyperboloid manifold with curvature c_out.

    Notes
    -----
    - Unlike HRC, f_t operates on the full point including the time component
    - f_t's output dimension determines the output spatial dimension
    - This is typically used for learnable linear transformations
    - The spatial scaling factor sqrt(c_in/c_out) ensures proper curvature transformation

    See Also
    --------
    hrc : Hyperbolic Regularization Component for spatial-only operations.
    HTCLinear : Module wrapper for htc with learnable linear transformation.

    References
    ----------
    Hypformer paper (citation to be added)

    Examples
    --------
    >>> import jax
    >>> import jax.numpy as jnp
    >>> from hyperbolix.nn_layers.hyperboloid_core import htc
    >>> from hyperbolix.manifolds import Hyperboloid
    >>>
    >>> # Create a point on the hyperboloid
    >>> manifold = Hyperboloid()
    >>> x = jnp.array([1.05, 0.1, -0.2, 0.15])
    >>> x = manifold.proj(x, c=1.0)
    >>>
    >>> # Define a linear transformation
    >>> W = jax.random.normal(jax.random.PRNGKey(0), (3, 4))
    >>> def linear(z):
    ...     return z @ W.T
    >>>
    >>> # Apply HTC
    >>> y = htc(x, linear, c_in=1.0, c_out=2.0)
    >>> y.shape
    (4,)  # (3 spatial + 1 time)
    """
    # f_t: (..., A_in) → (..., D_out) where A_in = in_dim+1
    out_D = f_t(x)

    # Curvature scaling sqrt(c_in/c_out) + time reconstruction (shared hrc/htc tail).
    return spatial_to_hyperboloid(out_D, c_in, c_out, eps)  # (..., D_out+1)

hyperbolix.nn_layers.hrc

hrc(
    x: Float[Array, "... dim_plus_1"],
    f_r: Callable[[Float[Array, ...]], Float[Array, ...]],
    c_in: float,
    c_out: float,
    eps: float = 1e-07,
) -> Float[Array, "... out_dim_plus_1"]

Hyperbolic Regularization Component.

Applies a Euclidean regularization/activation function f_r to the spatial components of hyperboloid points, then maps the result to the hyperboloid with curvature c_out.

Mathematical formula: space = sqrt(c_in/c_out) * f_r(x[..., 1:]) time = sqrt(||space||^2 + 1/c_out) output = [time, space]

When c_in = c_out = c, this reduces to: output = [sqrt(||f_r(x_s)||^2 + 1/c), f_r(x_s)] which is the pattern used by curvature-preserving hyperboloid activations.

Parameters:

Name	Type	Description	Default
x	Array of shape (..., dim+1)	Input point(s) on the hyperboloid manifold with curvature c_in. The first element is the time-like component, remaining are spatial.	required
f_r	Callable	Euclidean function to apply to spatial components. Can be any activation, normalization, dropout, etc. Takes spatial components and returns transformed spatial components (may change dimension).	required
c_in	float	Input curvature parameter (must be positive, c > 0).	required
c_out	float	Output curvature parameter (must be positive, c > 0).	required
eps	float	Small value for numerical stability (default: 1e-7).	1e-07

Returns:

Name	Type	Description
y	Array of shape (..., out_dim+1)	Output point(s) on the hyperboloid manifold with curvature c_out.

Notes

• f_r operates only on spatial components x[..., 1:], not the time component
• The time component is reconstructed using the hyperboloid constraint: -x₀² + ||x_rest||² = -1/c_out
• This avoids expensive exp/log maps while maintaining mathematical correctness
• The spatial scaling factor sqrt(c_in/c_out) ensures proper curvature transformation

See Also

htc : Hyperbolic Transformation Component for full-point operations.

References

Hypformer paper (citation to be added)

Examples:

>>> import jax.numpy as jnp
>>> from hyperbolix.nn_layers.hyperboloid_core import hrc
>>> from hyperbolix.manifolds import Hyperboloid
>>>
>>> # Create a point on the hyperboloid
>>> manifold = Hyperboloid()
>>> x = jnp.array([1.05, 0.1, -0.2, 0.15])
>>> x = manifold.proj(x, c=1.0)
>>>
>>> # Apply HRC with ReLU (curvature-preserving)
>>> y = hrc(x, jax.nn.relu, c_in=1.0, c_out=1.0)
>>>
>>> # Apply HRC with curvature change
>>> y = hrc(x, jax.nn.relu, c_in=1.0, c_out=2.0)
>>>
>>> # Custom activation
>>> def custom_act(z):
...     return jax.nn.gelu(z) * 0.5
>>> y = hrc(x, custom_act, c_in=1.0, c_out=0.5)

Source code in hyperbolix/nn_layers/hyperboloid_core.py

def hrc(
    x: Float[Array, "... dim_plus_1"],
    f_r: Callable[[Float[Array, "..."]], Float[Array, "..."]],
    c_in: float,
    c_out: float,
    eps: float = 1e-7,
) -> Float[Array, "... out_dim_plus_1"]:
    """Hyperbolic Regularization Component.

    Applies a Euclidean regularization/activation function f_r to the spatial
    components of hyperboloid points, then maps the result to the hyperboloid
    with curvature c_out.

    Mathematical formula:
        space = sqrt(c_in/c_out) * f_r(x[..., 1:])
        time  = sqrt(||space||^2 + 1/c_out)
        output = [time, space]

    When c_in = c_out = c, this reduces to:
        output = [sqrt(||f_r(x_s)||^2 + 1/c), f_r(x_s)]
    which is the pattern used by curvature-preserving hyperboloid activations.

    Parameters
    ----------
    x : Array of shape (..., dim+1)
        Input point(s) on the hyperboloid manifold with curvature c_in.
        The first element is the time-like component, remaining are spatial.
    f_r : Callable
        Euclidean function to apply to spatial components. Can be any activation,
        normalization, dropout, etc. Takes spatial components and returns
        transformed spatial components (may change dimension).
    c_in : float
        Input curvature parameter (must be positive, c > 0).
    c_out : float
        Output curvature parameter (must be positive, c > 0).
    eps : float, optional
        Small value for numerical stability (default: 1e-7).

    Returns
    -------
    y : Array of shape (..., out_dim+1)
        Output point(s) on the hyperboloid manifold with curvature c_out.

    Notes
    -----
    - f_r operates only on spatial components x[..., 1:], not the time component
    - The time component is reconstructed using the hyperboloid constraint:
      -x₀² + ||x_rest||² = -1/c_out
    - This avoids expensive exp/log maps while maintaining mathematical correctness
    - The spatial scaling factor sqrt(c_in/c_out) ensures proper curvature transformation

    See Also
    --------
    htc : Hyperbolic Transformation Component for full-point operations.

    References
    ----------
    Hypformer paper (citation to be added)

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> from hyperbolix.nn_layers.hyperboloid_core import hrc
    >>> from hyperbolix.manifolds import Hyperboloid
    >>>
    >>> # Create a point on the hyperboloid
    >>> manifold = Hyperboloid()
    >>> x = jnp.array([1.05, 0.1, -0.2, 0.15])
    >>> x = manifold.proj(x, c=1.0)
    >>>
    >>> # Apply HRC with ReLU (curvature-preserving)
    >>> y = hrc(x, jax.nn.relu, c_in=1.0, c_out=1.0)
    >>>
    >>> # Apply HRC with curvature change
    >>> y = hrc(x, jax.nn.relu, c_in=1.0, c_out=2.0)
    >>>
    >>> # Custom activation
    >>> def custom_act(z):
    ...     return jax.nn.gelu(z) * 0.5
    >>> y = hrc(x, custom_act, c_in=1.0, c_out=0.5)
    """
    x_space_D = x[..., 1:]  # (..., D) spatial components

    out_space_D = f_r(x_space_D)  # (..., D') — may change dim

    # Curvature scaling sqrt(c_in/c_out) + time reconstruction (shared hrc/htc tail).
    return spatial_to_hyperboloid(out_space_D, c_in, c_out, eps)  # (..., D'+1)

Point assembly & concatenation

Helpers that build valid hyperboloid points from spatial/tangent data and extract convolution patches.

hyperbolix.nn_layers.spatial_to_hyperboloid

spatial_to_hyperboloid(
    spatial: Float[Array, "... D"],
    c_in: float,
    c_out: float,
    eps: float = 1e-07,
) -> Float[Array, "... D_plus_1"]

Scale spatial components and reconstruct time to produce a hyperboloid point.

Extracts the common tail of hrc/htc: curvature scaling + time reconstruction via the hyperboloid constraint.

Parameters:

Name	Type	Description	Default
spatial	(Array, shape(..., D))	Spatial components (no time coordinate).	required
c_in	float	Source curvature (positive).	required
c_out	float	Target curvature (positive).	required
eps	float	Numerical stability floor (default: 1e-7).	1e-07

Returns:

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

Source code in hyperbolix/nn_layers/hyperboloid_core.py

def spatial_to_hyperboloid(
    spatial: Float[Array, "... D"],
    c_in: float,
    c_out: float,
    eps: float = 1e-7,
) -> Float[Array, "... D_plus_1"]:
    """Scale spatial components and reconstruct time to produce a hyperboloid point.

    Extracts the common tail of ``hrc``/``htc``: curvature scaling + time
    reconstruction via the hyperboloid constraint.

    Parameters
    ----------
    spatial : Array, shape (..., D)
        Spatial components (no time coordinate).
    c_in : float
        Source curvature (positive).
    c_out : float
        Target curvature (positive).
    eps : float, optional
        Numerical stability floor (default: 1e-7).

    Returns
    -------
    Array, shape (..., D+1)
        Points on the hyperboloid with curvature ``c_out``.
    """
    scale = jnp.sqrt(c_in / c_out)
    scaled_D = scale * spatial  # (..., D)

    norm_sq = jnp.sum(scaled_D**2, axis=-1)  # (...)
    x0 = jnp.sqrt(jnp.maximum(norm_sq + 1.0 / c_out, eps))  # (...)

    return jnp.concatenate([x0[..., None], scaled_D], axis=-1)  # (..., D+1)

hyperbolix.nn_layers.sinh_lift_to_hyperboloid

sinh_lift_to_hyperboloid(
    spatial_BO: Float[Array, "... O"],
    c: float,
    v_max: float,
    eps: float = 1e-07,
) -> Float[Array, "... O_plus_1"]

Lift spatial pre-activations to the hyperboloid via the element-wise sinh diffeomorphism.

Shared output map of HypLinearHyperboloidPLFC (Shi et al. 2026, point-to-hyperplane) and HypLinearHyperboloidBusemann (Chen et al. 2026, Thm. 4.2, point-to-horosphere): both compute a per-output score, then recover a hyperboloid point by ::

y_s = sinh(clip(√c · score, ±v_max)) / √c
y_t = √(‖y_s‖² + 1/c)

The clip is an output-side overflow guard: compute_mlr / the Busemann logit only bound an asinh/log argument, leaving the score ∝ ‖kernel_row‖ unbounded, and the sinh exponentiates it. Hard jnp.clip (not smooth_clamp) matches the Shi et al. 2026 reference and is value-identical for |√c·score| ≤ v_max — it only zeroes the gradient in the already-saturated tail. Callers must _assert_v_max_safe(v_max) so the bare jnp.sinh below cannot overflow float32.

Parameters:

Name	Type	Description	Default
spatial_BO	(Array, shape(..., O))	Per-output scores (the sinh argument, pre-scaling). O = out_dim - 1.	required
c	float	Curvature parameter (positive).	required
v_max	float	Hard clip bound on √c · spatial (output-side overflow guard).	required
eps	float	Numerical stability floor passed to :func:spatial_to_hyperboloid (default: 1e-7).	1e-07

Returns:

Type	Description
(Array, shape(..., O + 1))	Points on the hyperboloid with curvature c (time coordinate first).

Source code in hyperbolix/nn_layers/hyperboloid_core.py

def sinh_lift_to_hyperboloid(
    spatial_BO: Float[Array, "... O"],
    c: float,
    v_max: float,
    eps: float = 1e-7,
) -> Float[Array, "... O_plus_1"]:
    """Lift spatial pre-activations to the hyperboloid via the element-wise sinh diffeomorphism.

    Shared output map of ``HypLinearHyperboloidPLFC`` (Shi et al. 2026, point-to-hyperplane)
    and ``HypLinearHyperboloidBusemann`` (Chen et al. 2026, Thm. 4.2, point-to-horosphere):
    both compute a per-output score, then recover a hyperboloid point by ::

        y_s = sinh(clip(√c · score, ±v_max)) / √c
        y_t = √(‖y_s‖² + 1/c)

    The ``clip`` is an output-side overflow guard: ``compute_mlr`` / the Busemann logit only
    bound an ``asinh``/``log`` argument, leaving the score ∝ ``‖kernel_row‖`` unbounded, and
    the ``sinh`` exponentiates it. Hard ``jnp.clip`` (not ``smooth_clamp``) matches the Shi
    et al. 2026 reference and is value-identical for ``|√c·score| ≤ v_max`` — it only zeroes
    the gradient in the already-saturated tail. Callers must ``_assert_v_max_safe(v_max)`` so
    the bare ``jnp.sinh`` below cannot overflow float32.

    Parameters
    ----------
    spatial_BO : Array, shape (..., O)
        Per-output scores (the sinh argument, pre-scaling). ``O = out_dim - 1``.
    c : float
        Curvature parameter (positive).
    v_max : float
        Hard clip bound on ``√c · spatial`` (output-side overflow guard).
    eps : float, optional
        Numerical stability floor passed to :func:`spatial_to_hyperboloid` (default: 1e-7).

    Returns
    -------
    Array, shape (..., O+1)
        Points on the hyperboloid with curvature ``c`` (time coordinate first).
    """
    sqrt_c = jnp.sqrt(c)
    sinh_arg_BO = jnp.clip(sqrt_c * spatial_BO, -v_max, v_max)
    # Bare jnp.sinh: the argument is already bounded to ±v_max ≪ the math_utils.sinh overflow
    # clamp, so the safe variant would only re-clamp redundantly (callers assert v_max safety).
    res_rem_BO = jnp.sinh(sinh_arg_BO) / sqrt_c
    # Time reconstruction via the hyperboloid constraint (scale = 1 since c_in == c_out).
    return spatial_to_hyperboloid(res_rem_BO, c_in=c, c_out=c, eps=eps)

hyperbolix.nn_layers.build_spacelike_V

build_spacelike_V(
    U_IO: Float[Array, "I O"],
    b_O: Float[Array, O],
    c: float,
    eps: float = 1e-07,
) -> Float[Array, "Ai O"]

Build spacelike V matrix with Minkowski metric absorbed.

Constructs the spacelike weight vectors from Euclidean weights U and bias b via parallel transport (Eq. 12, Klis et al. 2026). The Minkowski metric signature diag(-1, +1, ..., +1) is absorbed into the time row so that the forward pass reduces to a single x @ V matmul.

Parameters:

Name	Type	Description	Default
U_IO	(Array, shape(I, O))	Euclidean weight matrix. I = in_spatial, O = out_spatial.	required
b_O	(Array, shape(O))	Bias per output neuron.	required
c	float	Curvature parameter (positive).	required
eps	float	Numerical stability floor for column norms (default: 1e-7).	1e-07

Returns:

Name	Type	Description
V_AiO	(Array, shape(I + 1, O))	Spacelike V matrix with negated time row (Minkowski metric absorbed). Columns are v^(i) with v_time_mink = -||w|| * sinh(arg).

References

Klis et al. "Fast and Geometrically Grounded Lorentz Neural Networks" (2026), Eq. 12.

Source code in hyperbolix/nn_layers/hyperboloid_core.py

def build_spacelike_V(
    U_IO: Float[Array, "I O"],
    b_O: Float[Array, "O"],
    c: float,
    eps: float = 1e-7,
) -> Float[Array, "Ai O"]:
    """Build spacelike V matrix with Minkowski metric absorbed.

    Constructs the spacelike weight vectors from Euclidean weights U and bias b
    via parallel transport (Eq. 12, Klis et al. 2026). The Minkowski metric
    signature diag(-1, +1, ..., +1) is absorbed into the time row so that the
    forward pass reduces to a single ``x @ V`` matmul.

    Parameters
    ----------
    U_IO : Array, shape (I, O)
        Euclidean weight matrix. I = in_spatial, O = out_spatial.
    b_O : Array, shape (O,)
        Bias per output neuron.
    c : float
        Curvature parameter (positive).
    eps : float, optional
        Numerical stability floor for column norms (default: 1e-7).

    Returns
    -------
    V_AiO : Array, shape (I+1, O)
        Spacelike V matrix with negated time row (Minkowski metric absorbed).
        Columns are v^(i) with ``v_time_mink = -||w|| * sinh(arg)``.

    References
    ----------
    Klis et al. "Fast and Geometrically Grounded Lorentz Neural Networks" (2026), Eq. 12.
    """
    # Dimension key: I=in_spatial, O=out_spatial, Ai=in_ambient (I+1)

    # Column norms of U: ||w^(i)||_E
    norm_sq_O = jnp.sum(U_IO**2, axis=0)  # (O,)
    norm_O = jnp.sqrt(norm_sq_O + eps)  # (O,) gradient-safe (never 0)

    # Smooth gate: 0 for zero-norm columns, ~1 for normal columns.
    # Ensures arg→0 smoothly when U column→0 (no weight → no bias transport).
    gate_O = norm_sq_O / (norm_sq_O + eps)  # (O,)

    # Argument for sinh/cosh: -sqrt(c) * b / ||w||, gated for zero columns
    arg_O = jnp.clip(-jnp.sqrt(c) * b_O * gate_O / norm_O, -100.0, 100.0)  # (O,)

    # Time row: negate to absorb Minkowski metric I_{1,D}
    # v_time = ||w|| * sinh(arg), v_time_mink = -v_time
    v_time_mink_O = -norm_O * safe_sinh(arg_O)  # (O,)

    # Spatial rows: w^(i) * cosh(arg)
    v_space_IO = U_IO * safe_cosh(arg_O)[None, :]  # (I, O)

    # Stack: [v_time_mink; v_space] -> (I+1, O) = (Ai, O)
    V_AiO = jnp.concatenate([v_time_mink_O[None, :], v_space_IO], axis=0)  # (Ai, O)

    return V_AiO

hyperbolix.nn_layers.extract_patches

extract_patches(
    x: Float[Array, "batch height width in_channels"],
    kernel_size: tuple[int, int],
    stride: tuple[int, int],
    padding: str,
    pad_mode: str = "edge",
    c: float | None = None,
) -> Float[
    Array,
    "batch out_height out_width kernel_h kernel_w in_channels",
]

Extract receptive-field patches for hyperboloid convolutions.

Shared by all hyperboloid conv layers (HypConv2DHyperboloid, HypConv2DHyperboloidFHNN, FGGConv2D, HypConv2DHyperboloidILNN). The patches are returned in point-major (kh, kw, C) order so that the downstream per-point ops (hcat / log_radius_concat) can consume them as (K, C) after flattening; this is why the channel-major output of jax.lax.conv_general_dilated_patches is transposed.

Padding is applied manually because the conv primitive zero-pads, and zero is not a point on the hyperboloid:

• pad_mode="edge" — replicate the border (default).
• pad_mode="origin" — fill with the manifold origin (sqrt(1/c), 0, ..., 0); c is required in this mode (matching the Shi et al. 2026 / Klis et al. 2026 reference padding).

Parameters:

Name	Type	Description	Default
x	(Array, shape(B, H, W, in_channels))	Input feature map of hyperboloid points (ambient channels, time first).	required
kernel_size	tuple[int, int]	(kernel_h, kernel_w).	required
stride	tuple[int, int]	(stride_h, stride_w).	required
padding	str	"SAME" or "VALID". Only "SAME" triggers manual padding.	required
pad_mode	str	"edge" or "origin" (default: "edge").	'edge'
c	float or None	Curvature, required only for pad_mode="origin" (default: None).	None

Returns:

Type	Description
(Array, shape(B, out_height, out_width, kernel_h, kernel_w, in_channels))	Receptive-field patches in point-major order.

Source code in hyperbolix/nn_layers/hyperboloid_core.py

def extract_patches(
    x: Float[Array, "batch height width in_channels"],
    kernel_size: tuple[int, int],
    stride: tuple[int, int],
    padding: str,
    pad_mode: str = "edge",
    c: float | None = None,
) -> Float[Array, "batch out_height out_width kernel_h kernel_w in_channels"]:
    """Extract receptive-field patches for hyperboloid convolutions.

    Shared by all hyperboloid conv layers (``HypConv2DHyperboloid``,
    ``HypConv2DHyperboloidFHNN``, ``FGGConv2D``, ``HypConv2DHyperboloidILNN``). The
    patches are returned in **point-major** ``(kh, kw, C)`` order so that the
    downstream per-point ops (``hcat`` / ``log_radius_concat``) can consume them as
    ``(K, C)`` after flattening; this is why the channel-major output of
    ``jax.lax.conv_general_dilated_patches`` is transposed.

    Padding is applied **manually** because the conv primitive zero-pads, and zero is
    not a point on the hyperboloid:

    - ``pad_mode="edge"``   — replicate the border (default).
    - ``pad_mode="origin"`` — fill with the manifold origin ``(sqrt(1/c), 0, ..., 0)``;
      ``c`` is required in this mode (matching the Shi et al. 2026 / Klis et al. 2026
      reference padding).

    Parameters
    ----------
    x : Array, shape (B, H, W, in_channels)
        Input feature map of hyperboloid points (ambient channels, time first).
    kernel_size : tuple[int, int]
        ``(kernel_h, kernel_w)``.
    stride : tuple[int, int]
        ``(stride_h, stride_w)``.
    padding : str
        ``"SAME"`` or ``"VALID"``. Only ``"SAME"`` triggers manual padding.
    pad_mode : str, optional
        ``"edge"`` or ``"origin"`` (default: ``"edge"``).
    c : float or None, optional
        Curvature, required only for ``pad_mode="origin"`` (default: None).

    Returns
    -------
    Array, shape (B, out_height, out_width, kernel_h, kernel_w, in_channels)
        Receptive-field patches in point-major order.
    """
    batch, height, width, in_channels = x.shape
    kernel_h, kernel_w = kernel_size
    stride_h, stride_w = stride

    # 1. Manual padding (the conv primitive's zero-pad is off-manifold)
    if padding == "SAME":
        out_height = (height + stride_h - 1) // stride_h
        out_width = (width + stride_w - 1) // stride_w
        pad_h = max((out_height - 1) * stride_h + kernel_h - height, 0)
        pad_w = max((out_width - 1) * stride_w + kernel_w - width, 0)
        pad_top = pad_h // 2
        pad_bottom = pad_h - pad_top
        pad_left = pad_w // 2
        pad_right = pad_w - pad_left

        if pad_mode == "origin":
            # Pad with manifold origin: (√(1/c), 0, ..., 0).
            # Buffer dtype follows x to avoid silent float64 promotion under x64.
            padded_h = height + pad_h
            padded_w = width + pad_w
            padded = jnp.zeros((batch, padded_h, padded_w, in_channels), dtype=x.dtype)
            padded = padded.at[..., 0].set(jnp.sqrt(1.0 / c))
            x = padded.at[:, pad_top : pad_top + height, pad_left : pad_left + width, :].set(x)
        else:  # edge
            x = jnp.pad(
                x,
                ((0, 0), (pad_top, pad_bottom), (pad_left, pad_right), (0, 0)),
                mode="edge",
            )

    # 2. Extract patches — output: (B, H, W, C*kh*kw), channel-major (C, kh, kw)
    patches_flat_BHW_CKhKw = jax.lax.conv_general_dilated_patches(
        lhs=x,
        filter_shape=(kernel_h, kernel_w),
        window_strides=(stride_h, stride_w),
        padding="VALID",
        dimension_numbers=("NHWC", "OIHW", "NHWC"),
    )

    # 3. Reshape to separate channels and kernel dims, then transpose to point-major.
    # conv_general_dilated_patches always emits channel-major (C, kh, kw) regardless of
    # the rhs spec letters, so the transpose to (kh, kw, C) is required (and is cheap).
    out_h, out_w = patches_flat_BHW_CKhKw.shape[1], patches_flat_BHW_CKhKw.shape[2]
    patches_BHWCkhkw = patches_flat_BHW_CKhKw.reshape(batch, out_h, out_w, in_channels, kernel_h, kernel_w)
    patches_BHWkhkwC = patches_BHWCkhkw.transpose(0, 1, 2, 4, 5, 3)  # move C last

    return patches_BHWkhkwC

Midpoints & residuals

hyperbolix.nn_layers.lorentz_midpoint

lorentz_midpoint(
    points: Float[Array, "... M A"],
    weights: Float[Array, "... N M"],
    c: float,
    eps: float = 1e-07,
) -> Float[Array, "... N A"]

Weighted Lorentzian midpoint over M points.

Generalises :func:lorentz_residual (which handles two points) to an arbitrary weighted combination used by full attention aggregation and multi-head averaging.

Formula (HELM, Chen et al. 2024): h = weights @ points (weighted sum) mu = h / (sqrt(c) * ||h||_L)

where ||h||_L = sqrt(-<h,h>_L) and <h,h>_L = -h_0^2 + ||h_s||^2.

Parameters:

Name	Type	Description	Default
points	(Array, shape(..., M, A))	Points on the hyperboloid with curvature c. A = d + 1.	required
weights	(Array, shape(..., N, M))	Combination weights (e.g. attention weights, uniform 1/M).	required
c	float	Curvature parameter (positive).	required
eps	float	Numerical stability floor (default: 1e-7).	1e-07

Returns:

Type	Description
(Array, shape(..., N, A))	Midpoints on the hyperboloid with curvature c.

Source code in hyperbolix/nn_layers/hyperboloid_core.py

def lorentz_midpoint(
    points: Float[Array, "... M A"],
    weights: Float[Array, "... N M"],
    c: float,
    eps: float = 1e-7,
) -> Float[Array, "... N A"]:
    """Weighted Lorentzian midpoint over M points.

    Generalises :func:`lorentz_residual` (which handles two points) to an
    arbitrary weighted combination used by full attention aggregation and
    multi-head averaging.

    Formula (HELM, Chen et al. 2024):
        ``h = weights @ points``  (weighted sum)
        ``mu = h / (sqrt(c) * ||h||_L)``

    where ``||h||_L = sqrt(-<h,h>_L)`` and ``<h,h>_L = -h_0^2 + ||h_s||^2``.

    Parameters
    ----------
    points : Array, shape (..., M, A)
        Points on the hyperboloid with curvature ``c``.  ``A = d + 1``.
    weights : Array, shape (..., N, M)
        Combination weights (e.g. attention weights, uniform ``1/M``).
    c : float
        Curvature parameter (positive).
    eps : float, optional
        Numerical stability floor (default: 1e-7).

    Returns
    -------
    Array, shape (..., N, A)
        Midpoints on the hyperboloid with curvature ``c``.
    """
    # h = sum_m w_{n,m} * points_m  →  (..., N, A)
    h_NA = jnp.einsum("...nm,...ma->...na", weights, points)

    # Minkowski squared norm: <h,h>_L = -h_0^2 + ||h_s||^2  (should be < 0)
    mink_1 = -(h_NA[..., 0:1] ** 2) + jnp.sum(h_NA[..., 1:] ** 2, axis=-1, keepdims=True)  # (..., N, 1)
    denom_1 = jnp.sqrt(jnp.maximum(c * jnp.abs(mink_1), eps))  # (..., N, 1)

    return h_NA / denom_1  # (..., N, A)

hyperbolix.nn_layers.lorentz_residual

lorentz_residual(
    x: Float[Array, "... dim_plus_1"],
    y: Float[Array, "... dim_plus_1"],
    w_y: float | Float[Array, ""],
    c: float,
    eps: float = 1e-07,
) -> Float[Array, "... dim_plus_1"]

Lorentzian residual connection (LResNet, He et al. 2025).

Computes the weighted Lorentzian midpoint of x and y, projecting back to the hyperboloid:

ave = x + w_y * y
result = ave / sqrt(c * |<ave, ave>_L|)

where _L = -a_0^2 + ||a_s||^2 is the Minkowski inner product.

.. warning:: w_y must be non-negative. For x, y on the upper sheet, any conic combination x + w_y * y with w_y >= 0 stays future-directed timelike, so the normalization returns a valid hyperboloid point. For w_y < 0 the combination can turn spacelike (<ave, ave>_L > 0, roughly w_y < -1 for nearby points) or land on the lower sheet (ave_0 < 0) — the abs() in the normalizer then converts the geometry violation into a "valid-looking" but wrong output instead of raising. This is why callers must not expose w_y as an unconstrained learnable parameter.

Parameters:

Name	Type	Description	Default
x	(Array, shape(..., d + 1))	Points on hyperboloid with curvature c.	required
y	(Array, shape(..., d + 1))	Points on hyperboloid with curvature c (to be added with weight w_y).	required
w_y	float or scalar Array	Weight for the y contribution. Must be >= 0 (see warning above).	required
c	float	Curvature parameter (positive, c > 0).	required
eps	float	Numerical stability floor (default: 1e-7).	1e-07

Returns:

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

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. (Also adopted as the residual connection in HELM, Chen et al. 2024, Eq. 2.)

Source code in hyperbolix/nn_layers/hyperboloid_core.py

def lorentz_residual(
    x: Float[Array, "... dim_plus_1"],
    y: Float[Array, "... dim_plus_1"],
    w_y: float | Float[Array, ""],
    c: float,
    eps: float = 1e-7,
) -> Float[Array, "... dim_plus_1"]:
    """Lorentzian residual connection (LResNet, He et al. 2025).

    Computes the weighted Lorentzian midpoint of x and y, projecting back
    to the hyperboloid:

        ave = x + w_y * y
        result = ave / sqrt(c * |<ave, ave>_L|)

    where <a, a>_L = -a_0^2 + ||a_s||^2 is the Minkowski inner product.

    .. warning::
        ``w_y`` must be **non-negative**. For x, y on the upper sheet, any conic
        combination ``x + w_y * y`` with ``w_y >= 0`` stays future-directed
        timelike, so the normalization returns a valid hyperboloid point. For
        ``w_y < 0`` the combination can turn spacelike (``<ave, ave>_L > 0``,
        roughly ``w_y < -1`` for nearby points) or land on the lower sheet
        (``ave_0 < 0``) — the ``abs()`` in the normalizer then converts the
        geometry violation into a "valid-looking" but wrong output instead of
        raising. This is why callers must not expose ``w_y`` as an
        unconstrained learnable parameter.

    Parameters
    ----------
    x : Array, shape (..., d+1)
        Points on hyperboloid with curvature c.
    y : Array, shape (..., d+1)
        Points on hyperboloid with curvature c (to be added with weight w_y).
    w_y : float or scalar Array
        Weight for the y contribution. Must be >= 0 (see warning above).
    c : float
        Curvature parameter (positive, c > 0).
    eps : float, optional
        Numerical stability floor (default: 1e-7).

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

    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.
    (Also adopted as the residual connection in HELM, Chen et al. 2024, Eq. 2.)
    """
    ave_A = x + w_y * y  # (..., A) where A = d+1
    # Minkowski inner: -ave_0^2 + ||ave_s||^2
    mink_1 = -(ave_A[..., 0:1] ** 2) + jnp.sum(ave_A[..., 1:] ** 2, axis=-1, keepdims=True)  # (..., 1)
    denom_1 = jnp.sqrt(jnp.maximum(c * jnp.abs(mink_1), eps))  # (..., 1)
    return ave_A / denom_1  # (..., A)

hyperbolix.nn_layers.lorentz_scale

lorentz_scale(
    m: Float[Array, "... dim_plus_1"],
    gamma: float | Float[Array, ""],
    c: float,
    eps: float = 1e-07,
) -> Float[Array, "... dim_plus_1"]

Klein-geodesic output rescaling of a hyperboloid point (LResNet, Eq. 10).

Scales the spatial part by gamma and recomputes the time component so the result stays on the hyperboloid::

m_s' = gamma * m_s
m_t' = sqrt(||m_s'||^2 + 1/c)

The isometry to the Klein model is phi_K(x) = x_s / x_t and Klein geodesics are Euclidean lines through the origin, so this simply slides m along the geodesic ray from the origin -- changing the Euclidean / geodesic norm of the output while keeping it on the manifold. It restores expressiveness the bare Lorentzian residual lacks: :func:lorentz_residual's output is bounded by the larger of its two inputs, whereas gamma > 1 can push the norm beyond both.

Unlike lorentz_residual's w_y, any real gamma yields a valid hyperboloid point (m_t' > 0 always), so a learnable gamma is safe. gamma = 1 is the identity (for an on-manifold m, recomputing the time coordinate returns it unchanged); the LResNet reference uses a fixed gamma = 2 for its CIFAR configurations.

Parameters:

Name	Type	Description	Default
m	(Array, shape(..., d + 1))	Points on the hyperboloid with curvature c.	required
gamma	float or scalar Array	Scaling constant. Positive values slide toward (< 1) or away from (> 1) the origin along the geodesic ray.	required
c	float	Curvature parameter (positive, c > 0).	required
eps	float	Numerical stability floor (default: 1e-7).	1e-07

Returns:

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

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. Eq. (10) (optional output scaling). See :func:lorentz_residual for the residual this complements.

Source code in hyperbolix/nn_layers/hyperboloid_core.py

def lorentz_scale(
    m: Float[Array, "... dim_plus_1"],
    gamma: float | Float[Array, ""],
    c: float,
    eps: float = 1e-7,
) -> Float[Array, "... dim_plus_1"]:
    """Klein-geodesic output rescaling of a hyperboloid point (LResNet, Eq. 10).

    Scales the spatial part by ``gamma`` and recomputes the time component so the
    result stays on the hyperboloid::

        m_s' = gamma * m_s
        m_t' = sqrt(||m_s'||^2 + 1/c)

    The isometry to the Klein model is ``phi_K(x) = x_s / x_t`` and Klein geodesics
    are Euclidean lines through the origin, so this simply slides ``m`` along the
    geodesic ray from the origin -- changing the Euclidean / geodesic norm of the
    output while keeping it on the manifold. It restores expressiveness the bare
    Lorentzian residual lacks: :func:`lorentz_residual`'s output is bounded by the
    larger of its two inputs, whereas ``gamma > 1`` can push the norm beyond both.

    Unlike ``lorentz_residual``'s ``w_y``, *any* real ``gamma`` yields a valid
    hyperboloid point (``m_t' > 0`` always), so a learnable ``gamma`` is safe.
    ``gamma = 1`` is the identity (for an on-manifold ``m``, recomputing the time
    coordinate returns it unchanged); the LResNet reference uses a fixed
    ``gamma = 2`` for its CIFAR configurations.

    Parameters
    ----------
    m : Array, shape (..., d+1)
        Points on the hyperboloid with curvature ``c``.
    gamma : float or scalar Array
        Scaling constant. Positive values slide toward (``< 1``) or away from
        (``> 1``) the origin along the geodesic ray.
    c : float
        Curvature parameter (positive, c > 0).
    eps : float, optional
        Numerical stability floor (default: 1e-7).

    Returns
    -------
    Array, shape (..., d+1)
        Rescaled points on the hyperboloid with curvature ``c``.

    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.
    Eq. (10) (optional output scaling). See :func:`lorentz_residual` for the residual
    this complements.
    """
    spatial = m[..., 1:]  # (..., d)
    # Reuse the scale + time-reconstruction tail; c_in == c_out makes the curvature
    # factor sqrt(c/c) = 1, leaving a pure spatial gamma-scale.
    return spatial_to_hyperboloid(gamma * spatial, c_in=c, c_out=c, eps=eps)

Attention focus transform

hyperbolix.nn_layers.focus_transform

focus_transform(
    x_D: Float[Array, "... D"],
    temperature: Float[Array, ""],
    power: float,
    eps: float = 1e-07,
) -> Float[Array, "... D"]

Norm-preserving focus function (Eq 19, Hypformer).

Applies temperature-scaled ReLU followed by element-wise power sharpening while preserving the original norm.

Parameters:

Name	Type	Description	Default
x_D	(Array, shape(..., D))	Input spatial features.	required
temperature	scalar Array	Learnable temperature parameter.	required
power	float	Sharpening exponent (p > 1 concentrates mass).	required
eps	float	Numerical stability floor (default: 1e-7).	1e-07

Returns:

Type	Description
(Array, shape(..., D))	Focus-transformed features with ||output|| ≈ ||relu(x)/|t|||.

Source code in hyperbolix/nn_layers/hyperboloid_attention.py

def focus_transform(
    x_D: Float[Array, "... D"],
    temperature: Float[Array, ""],
    power: float,
    eps: float = 1e-7,
) -> Float[Array, "... D"]:
    """Norm-preserving focus function (Eq 19, Hypformer).

    Applies temperature-scaled ReLU followed by element-wise power sharpening
    while preserving the original norm.

    Parameters
    ----------
    x_D : Array, shape (..., D)
        Input spatial features.
    temperature : scalar Array
        Learnable temperature parameter.
    power : float
        Sharpening exponent (``p > 1`` concentrates mass).
    eps : float, optional
        Numerical stability floor (default: 1e-7).

    Returns
    -------
    Array, shape (..., D)
        Focus-transformed features with ``||output|| ≈ ||relu(x)/|t|||``.
    """
    # Temperature-scaled ReLU
    scaled_relu_D = (jax.nn.relu(x_D) + eps) / (jnp.abs(temperature) + eps)  # (..., D)

    # Element-wise power sharpening
    sharpened_D = scaled_relu_D**power  # (..., D)

    # Norm-preserving rescaling
    norm_scaled = jnp.sqrt(jnp.sum(scaled_relu_D**2, axis=-1, keepdims=True) + eps)  # (..., 1)
    norm_sharpened = jnp.sqrt(jnp.sum(sharpened_D**2, axis=-1, keepdims=True) + eps)  # (..., 1)

    return (norm_scaled / norm_sharpened) * sharpened_D  # (..., D)

Poincaré statistics & midpoints

Reductions used by PoincareBatchNorm2D and the Poincaré VQ codebook. frechet_variance is manifold-agnostic (mean squared geodesic distance); poincare_weighted_midpoint is the GGBall weighted gyromidpoint (the Poincaré analog of lorentz_midpoint).

hyperbolix.nn_layers.poincare_midpoint

poincare_midpoint(
    x_NC: Float[Array, "N C"],
    manifold: Poincare,
    c: float,
    eps: float = 1e-06,
) -> Float[Array, C]

Compute the (Einstein) gyromidpoint of Poincaré ball points.

Delegates to :func:poincare_weighted_midpoint (GGBall Eq. 41) with uniform weights::

mu = (1/2) ⊗_c [ Σ λ(x_n)·x_n / Σ (λ(x_n) - 1) ]

NOTE: an earlier version computed Σ(λ²·x) / Σ(λ²), which is not a midpoint (it is not equidistant for two points and is biased toward points near the ball boundary, where λ explodes).

Parameters:

Name	Type	Description	Default
x_NC	(Array, shape(N, C))	Points on the Poincaré ball.	required
manifold	Poincare	Poincaré manifold instance.	required
c	float	Curvature (positive).	required
eps	float	Numerical stability floor (default: 1e-6).	1e-06

Returns:

Type	Description
(Array, shape(C))	Gyromidpoint on the Poincaré ball.

Source code in hyperbolix/nn_layers/poincare_batchnorm.py

def poincare_midpoint(
    x_NC: Float[Array, "N C"],
    manifold: Poincare,
    c: float,
    eps: float = 1e-6,
) -> Float[Array, "C"]:
    """Compute the (Einstein) gyromidpoint of Poincaré ball points.

    Delegates to :func:`poincare_weighted_midpoint` (GGBall Eq. 41) with uniform
    weights::

        mu = (1/2) ⊗_c [ Σ λ(x_n)·x_n / Σ (λ(x_n) - 1) ]

    NOTE: an earlier version computed ``Σ(λ²·x) / Σ(λ²)``, which is *not* a
    midpoint (it is not equidistant for two points and is biased toward points
    near the ball boundary, where λ explodes).

    Parameters
    ----------
    x_NC : Array, shape (N, C)
        Points on the Poincaré ball.
    manifold : Poincare
        Poincaré manifold instance.
    c : float
        Curvature (positive).
    eps : float
        Numerical stability floor (default: 1e-6).

    Returns
    -------
    Array, shape (C,)
        Gyromidpoint on the Poincaré ball.
    """
    weights_1N = jnp.ones((1, x_NC.shape[0]), dtype=x_NC.dtype)  # (1, N) uniform
    return poincare_weighted_midpoint(x_NC, weights_1N, manifold, c, eps)[0]

hyperbolix.nn_layers.poincare_weighted_midpoint

poincare_weighted_midpoint(
    points_MC: Float[Array, "M C"],
    weights_NM: Float[Array, "N M"],
    manifold: Poincare,
    c: float,
    eps: float = 1e-06,
) -> Float[Array, "N C"]

Weighted Poincaré gyromidpoint over M points (GGBall Eq. 41).

Generalises :func:poincare_midpoint (the unweighted Einstein midpoint of N points) to N weighted midpoints, mirroring the hyperboloid :func:hyperbolix.nn_layers.lorentz_midpoint. For weight-vector w_n over the M points::

mu_n = (1/2) ⊗_c [ Σ_m w_nm · λ(x_m) · x_m  /  Σ_m w_nm · (λ(x_m) - 1) ]

where λ(x) = 2 / (1 - c‖x‖²) is the conformal factor and ⊗_c is Möbius scalar multiplication. The ½ ⊗_c half-scaling is exactly what turns the conformal-weighted Euclidean average into the gyro-midpoint on the ball (for N copies of a single point it returns that point — see the algebra in the tests).

This is the codebook-centroid operator of the hyperbolic-EMA VQ update (GGBall, Bu et al. 2026): pass the assigned encoder points as points and the transposed one-hot assignment matrix (num_codes, N) as weights to get one midpoint per code.

Parameters:

Name	Type	Description	Default
points_MC	(Array, shape(M, C))	Points on the Poincaré ball with curvature c.	required
weights_NM	(Array, shape(N, M))	Combination weights, one row per output midpoint. Typically non-negative (assignment / attention weights). Because λ(x) ≥ 2 the denominator Σ_m w_nm (λ-1) is strictly positive whenever a row carries any mass; an all-zero row maps to the origin.	required
manifold	Poincare	Poincaré manifold instance (supplies conformal_factor, scalar_mul, proj).	required
c	float	Curvature (positive).	required
eps	float	Denominator floor guarding the empty-row 0/0 (default: 1e-6).	1e-06

Returns:

Type	Description
(Array, shape(N, C))	Weighted gyromidpoints on the Poincaré ball.

Source code in hyperbolix/nn_layers/poincare_batchnorm.py

def poincare_weighted_midpoint(
    points_MC: Float[Array, "M C"],
    weights_NM: Float[Array, "N M"],
    manifold: Poincare,
    c: float,
    eps: float = 1e-6,
) -> Float[Array, "N C"]:
    """Weighted Poincaré gyromidpoint over M points (GGBall Eq. 41).

    Generalises :func:`poincare_midpoint` (the unweighted Einstein midpoint of N
    points) to N *weighted* midpoints, mirroring the hyperboloid
    :func:`hyperbolix.nn_layers.lorentz_midpoint`. For weight-vector ``w_n`` over
    the M points::

        mu_n = (1/2) ⊗_c [ Σ_m w_nm · λ(x_m) · x_m  /  Σ_m w_nm · (λ(x_m) - 1) ]

    where ``λ(x) = 2 / (1 - c‖x‖²)`` is the conformal factor and ``⊗_c`` is
    Möbius scalar multiplication. The ``½ ⊗_c`` half-scaling is exactly what turns
    the conformal-weighted Euclidean average into the gyro-midpoint *on* the ball
    (for N copies of a single point it returns that point — see the algebra in
    the tests).

    This is the codebook-centroid operator of the hyperbolic-EMA VQ update
    (GGBall, Bu et al. 2026): pass the assigned encoder points as ``points`` and
    the transposed one-hot assignment matrix ``(num_codes, N)`` as ``weights`` to
    get one midpoint per code.

    Parameters
    ----------
    points_MC : Array, shape (M, C)
        Points on the Poincaré ball with curvature ``c``.
    weights_NM : Array, shape (N, M)
        Combination weights, one row per output midpoint. Typically non-negative
        (assignment / attention weights). Because ``λ(x) ≥ 2`` the denominator
        ``Σ_m w_nm (λ-1)`` is strictly positive whenever a row carries any mass;
        an all-zero row maps to the origin.
    manifold : Poincare
        Poincaré manifold instance (supplies ``conformal_factor``, ``scalar_mul``,
        ``proj``).
    c : float
        Curvature (positive).
    eps : float
        Denominator floor guarding the empty-row 0/0 (default: 1e-6).

    Returns
    -------
    Array, shape (N, C)
        Weighted gyromidpoints on the Poincaré ball.
    """
    # lambda per point — conformal_factor is batched and returns (M, 1).
    lambda_M = manifold.conformal_factor(points_MC, c)[:, 0]  # (M,)

    # Eq. 41 numerator / denominator, contracted over the M axis:
    #   numerator_n = Σ_m w_nm · λ(x_m) · x_m   (N, C)
    #   denom_n     = Σ_m w_nm · (λ(x_m) - 1)   (N,)
    numerator_NC = jnp.einsum("nm,mc->nc", weights_NM, lambda_M[:, None] * points_MC)  # (N, C)
    denom_N = jnp.einsum("nm,m->n", weights_NM, lambda_M - 1.0)  # (N,)
    # Empty row (all-zero weights) -> denom 0 and numerator 0; map to origin
    # rather than 0/0. λ ≥ 2 keeps denom > 0 for any row with mass.
    denom_safe_N = jnp.where(jnp.abs(denom_N) < eps, 1.0, denom_N)  # (N,)
    inner_NC = numerator_NC / denom_safe_N[:, None]  # (N, C) — argument of ½ ⊗_c

    # Boundary guard before the Möbius half-scaling (keeps scalar_mul's atanh in domain).
    inner_NC = jax.vmap(manifold.proj, in_axes=(0, None))(inner_NC, c)

    # mu_n = ½ ⊗_c inner_n  (single-point scalar_mul, vmapped over N).
    mu_NC = jax.vmap(lambda v: manifold.scalar_mul(0.5, v, c))(inner_NC)  # (N, C)
    return jax.vmap(manifold.proj, in_axes=(0, None))(mu_NC, c)

hyperbolix.nn_layers.frechet_variance

frechet_variance(
    x_NC: Float[Array, "N C"],
    mean_C: Float[Array, C],
    manifold: Poincare,
    c: float,
) -> Float[Array, ""]

Compute Fréchet variance: mean squared geodesic distance to mean.

Parameters:

Name	Type	Description	Default
x_NC	(Array, shape(N, C))	Points on the Poincaré ball.	required
mean_C	(Array, shape(C))	Mean point on the Poincaré ball.	required
manifold	Poincare	Poincaré manifold instance.	required
c	float	Curvature (positive).	required

Returns:

Type	Description
(Array, scalar)	Mean of squared geodesic distances.

Source code in hyperbolix/nn_layers/poincare_batchnorm.py

def frechet_variance(
    x_NC: Float[Array, "N C"],
    mean_C: Float[Array, "C"],
    manifold: Poincare,
    c: float,
) -> Float[Array, ""]:
    """Compute Fréchet variance: mean squared geodesic distance to mean.

    Parameters
    ----------
    x_NC : Array, shape (N, C)
        Points on the Poincaré ball.
    mean_C : Array, shape (C,)
        Mean point on the Poincaré ball.
    manifold : Poincare
        Poincaré manifold instance.
    c : float
        Curvature (positive).

    Returns
    -------
    Array, scalar
        Mean of squared geodesic distances.
    """
    # dist per point: (N,)
    dists_N = jax.vmap(manifold.dist, in_axes=(0, None, None))(x_NC, mean_C, c)
    return jnp.mean(dists_N**2)

Example

import jax, jax.numpy as jnp
from hyperbolix.nn_layers import lorentz_residual
from hyperbolix.manifolds import Hyperboloid

hyperboloid, c, d = Hyperboloid(), 1.0, 6
x = jax.random.normal(jax.random.PRNGKey(0), (8, d + 1))
y = jax.random.normal(jax.random.PRNGKey(1), (8, d + 1))
x = jax.vmap(hyperboloid.proj, in_axes=(0, None))(x, c)
y = jax.vmap(hyperboloid.proj, in_axes=(0, None))(y, c)
result = lorentz_residual(x, y, w_y=0.3, c=c)  # weighted Lorentzian midpoint (skip connection)
print(result.shape)  # (8, 7)