Normalization

All hyperbolic normalization layers in one place — previously these were scattered across the convolution and Hypformer sections. For which normalizer to use between which layers, see the NN Layers guide.

Channel convention

HRC and gyro norms take the spatial dimension num_features = d (not ambient). PoincareBatchNorm2D also takes spatial d (Poincaré has no time component).

Poincaré Batch Normalization

PoincareBatchNorm2D operates in tangent space (matching conv-layer I/O), mapping to the manifold internally for the geometric operations (Einstein midpoint, Fréchet variance, parallel transport, variance rescaling). Use between Poincaré conv layers following the ResNet pattern conv → bn → relu → conv → bn → skip. The midpoint / variance helpers it builds on live in Primitives.

hyperbolix.nn_layers.PoincareBatchNorm2D

PoincareBatchNorm2D(
    manifold_module: Poincare,
    num_features: int,
    *,
    use_running_average: bool = False,
    momentum: float = 0.9,
    eps: float = 1e-06,
    param_dtype: DTypeLike = jnp.float32,
)

Bases: Module

Poincaré batch normalization for 2D feature maps.

Operates in tangent space (matching conv layer I/O). Internally maps to the manifold for geometric operations (midpoint, parallel transport, variance rescaling), then maps back to tangent space.

Follows nnx.BatchNorm interface: use_running_average is a constructor parameter overridable at call time.

Parameters:

Name	Type	Description	Default
manifold_module	Poincare	Poincaré manifold instance.	required
num_features	int	Number of channels (Poincaré ball dimension per pixel).	required
use_running_average	bool	If True, use running statistics instead of batch statistics (default: False). Overridable at call time.	False
momentum	float	EMA momentum for running statistics (default: 0.9).	0.9
eps	float	Numerical stability floor (default: 1e-6).	1e-06
param_dtype	DTypeLike	Storage dtype of the learnable parameters and running statistics (default: jnp.float32). Compute precision of manifold operations is set by manifold.dtype.	float32

Notes

The learned variance self.var is an unconstrained scalar used as sqrt(var / (batch_var + eps)) — if gradient descent drives it negative, the sqrt produces NaN. This matches the reference (van Spengler's PoincareBatchNorm uses an unconstrained nn.Parameter under the same sqrt), so we keep the parameterization; in practice the init at 1.0 and typical learning rates keep it positive. If you observe NaNs originating here, reparameterize via softplus in your model or clamp self.var after optimizer steps.

References

van Spengler et al. "Poincaré ResNet." ICML 2023.

Source code in hyperbolix/nn_layers/poincare_batchnorm.py

def __init__(
    self,
    manifold_module: Poincare,
    num_features: int,
    *,
    use_running_average: bool = False,
    momentum: float = 0.9,
    eps: float = 1e-6,
    param_dtype: DTypeLike = jnp.float32,
):
    validate_poincare_manifold(
        manifold_module,
        required_methods=("expmap_0", "logmap_0", "logmap", "expmap", "ptransp", "proj", "dist", "conformal_factor"),
    )
    self.manifold = manifold_module
    self.num_features = num_features
    self.use_running_average = use_running_average
    self.momentum = momentum
    self.eps = eps

    # Learnable parameters and running stats (tangent space). Storage is
    # pinned to param_dtype; manifold operations re-promote to the
    # manifold's compute dtype at their boundaries, and the EMA updates in
    # __call__ cast back to the stat dtype at store time.
    self.mean = nnx.Param(jnp.zeros((num_features,), dtype=param_dtype))  # learned mean
    self.var = nnx.Param(jnp.ones((), dtype=param_dtype))  # learned variance (scalar)

    # Running statistics (tangent space)
    self.running_mean = nnx.BatchStat(jnp.zeros((num_features,), dtype=param_dtype))
    self.running_var = nnx.BatchStat(jnp.ones((), dtype=param_dtype))

__call__

__call__(
    x_BHWC: Float[Array, "B H W C"],
    c: float = 1.0,
    use_running_average: bool | None = None,
) -> Float[Array, "B H W C"]

Apply Poincaré batch normalization.

Parameters:

Name	Type	Description	Default
x_BHWC	(Array, shape(B, H, W, C))	Tangent-space input features.	required
c	float	Curvature (positive, default: 1.0).	1.0
use_running_average	bool or None	Override constructor setting. None uses constructor value.	None

Returns:

Type	Description
(Array, shape(B, H, W, C))	Normalized tangent-space output.

Source code in hyperbolix/nn_layers/poincare_batchnorm.py

def __call__(
    self,
    x_BHWC: Float[Array, "B H W C"],
    c: float = 1.0,
    use_running_average: bool | None = None,
) -> Float[Array, "B H W C"]:
    """Apply Poincaré batch normalization.

    Parameters
    ----------
    x_BHWC : Array, shape (B, H, W, C)
        Tangent-space input features.
    c : float
        Curvature (positive, default: 1.0).
    use_running_average : bool or None
        Override constructor setting. None uses constructor value.

    Returns
    -------
    Array, shape (B, H, W, C)
        Normalized tangent-space output.
    """
    # Resolve use_running_average: call-time > constructor
    if use_running_average is None:
        use_running_average = self.use_running_average

    B, H, W, C = x_BHWC.shape

    # Flatten (B, H, W, C) → (N, C)
    x_NC = x_BHWC.reshape(-1, C)  # (N, C)

    # --- Map to manifold ---
    # x_NC are tangent vectors; map to Poincaré ball
    x_manifold_NC = jax.vmap(self.manifold.expmap_0, in_axes=(0, None))(x_NC, c)  # (N, C)

    if use_running_average:
        # Use running stats (eval mode)
        batch_mean_tangent_C = self.running_mean[...]  # (C,)
        batch_var = self.running_var[...]  # scalar
    else:
        # Compute batch stats on manifold
        batch_midpoint_C = poincare_midpoint(x_manifold_NC, self.manifold, c, self.eps)  # (C,)
        batch_var = frechet_variance(x_manifold_NC, batch_midpoint_C, self.manifold, c)  # scalar

        # Map midpoint back to tangent space for running stat storage
        batch_mean_tangent_C = self.manifold.logmap_0(batch_midpoint_C, c)  # (C,)

        # Update running statistics (EMA, no gradient flow). The batch
        # stats come from manifold ops (manifold.dtype), so cast back to
        # the stat storage dtype at store time.
        new_running_mean_C = self.momentum * self.running_mean[...] + (1.0 - self.momentum) * batch_mean_tangent_C
        self.running_mean[...] = jax.lax.stop_gradient(new_running_mean_C).astype(self.running_mean[...].dtype)
        new_running_var = self.momentum * self.running_var[...] + (1.0 - self.momentum) * batch_var
        self.running_var[...] = jax.lax.stop_gradient(new_running_var).astype(self.running_var[...].dtype)

    # Get the manifold-space mean to use for geometric operations
    # (either from batch or from running stats)
    if use_running_average:
        input_mean_C = self.manifold.expmap_0(batch_mean_tangent_C, c)  # (C,)
        current_var = batch_var
    else:
        input_mean_C = batch_midpoint_C  # already computed above  # (C,)
        current_var = batch_var

    # Learned mean on manifold
    learned_mean_C = self.manifold.expmap_0(self.mean[...], c)  # (C,)

    # --- Step 4: logmap all points at input mean ---
    v_NC = jax.vmap(self.manifold.logmap, in_axes=(0, None, None))(x_manifold_NC, input_mean_C, c)  # (N, C)

    # --- Step 5: parallel transport from input mean to learned mean ---
    v_NC = jax.vmap(self.manifold.ptransp, in_axes=(0, None, None, None))(v_NC, input_mean_C, learned_mean_C, c)  # (N, C)

    # --- Step 6: rescale by sqrt(learned_var / (batch_var + eps)) ---
    scale = jnp.sqrt(self.var[...] / (current_var + self.eps))  # scalar
    v_NC = v_NC * scale  # (N, C)

    # --- Step 7: expmap at learned mean ---
    out_NC = jax.vmap(self.manifold.expmap, in_axes=(0, None, None))(v_NC, learned_mean_C, c)  # (N, C)

    # --- Step 8: logmap_0 back to tangent space ---
    out_NC = jax.vmap(self.manifold.logmap_0, in_axes=(0, None))(out_NC, c)  # (N, C)

    # Reshape back to (B, H, W, C)
    return out_NC.reshape(B, H, W, C)

Gyro Normalization (Hyperboloid, Proper Velocity & Poincaré)

Intrinsic normalization that operates directly on manifold points via gyrovector operations (addition, scalar_mul) — no tangent-space round trip for the affine part. Batch normalization is provided for the Lorentz/Hyperboloid and Proper Velocity models (the Poincaré ball's batch normalizer is the tangent-space PoincareBatchNorm2D above); the per-sample radial RMSNorm covers all three.

• Gyrogroup Batch Normalization (HyperboloidGyroBatchNorm, ProperVelocityGyroBatchNorm) — a port of GyroBN (Chen et al., ICLR 2024 / 2025). Centers by gyro-translating with the inverse batch mean, scales by the inverse Fréchet standard deviation, biases with a learned manifold point, and keeps running statistics for evaluation (use_running_average). The Hyperboloid batch mean is the closed-form Lorentz centroid; the PV mean is the closed-form log-Euclidean mean. Use for faithful hyperbolic ResNets.
• Gyro radial RMSNorm (HyperboloidGyroRMSNorm, ProperVelocityGyroRMSNorm, PoincareGyroRMSNorm) — a per-sample, batch-independent normalizer (no running statistics, identical in train and eval, valid at batch size 1 — the properties RL workloads want). Each point's geodesic radius is rescaled to a learned target gamma via a single gyro scalar-multiplication, with optional gyro-bias (use_bias). The manifold analog of RMSNorm: normalizes magnitude (hierarchy depth) while preserving direction. Möbius scalar_mul (Poincaré) and Lorentz/PV scalar_mul scale geodesic radius identically, so the same layer body serves all three models.

hyperbolix.nn_layers.HyperboloidGyroBatchNorm

HyperboloidGyroBatchNorm(
    manifold_module: Hyperboloid,
    num_features: int,
    **kwargs,
)

Bases: GyroBatchNormBase

GyroBN for the Hyperboloid (Lorentz) model.

Inputs are ambient (..., D+1) hyperboloid points; num_features is the spatial dimension D. The batch mean is the closed-form Lorentz centroid (HELM, Chen et al. 2024) via :func:lorentz_midpoint — exact and JIT-friendly, matching the estimator the ILNN GyroBN reference uses in practice.

Source code in hyperbolix/nn_layers/gyro_normalization.py

def __init__(self, manifold_module: Hyperboloid, num_features: int, **kwargs):
    validate_hyperboloid_manifold(manifold_module, required_methods=_GYRO_BN_METHODS)
    super().__init__(manifold_module, num_features, **kwargs)

hyperbolix.nn_layers.ProperVelocityGyroBatchNorm

ProperVelocityGyroBatchNorm(
    manifold_module: ProperVelocity,
    num_features: int,
    **kwargs,
)

Bases: GyroBatchNormBase

GyroBN for the Proper Velocity (PV) model.

Inputs are (..., D) PV points; num_features = D. PV has no closed-form centroid, so the batch mean is the closed-form log-Euclidean mean expmap_0(mean_i logmap_0(x_i)) (the GyroBN reference's use_euclid_stats mode): no iteration, fully vmap/JIT-clean.

Source code in hyperbolix/nn_layers/gyro_normalization.py

def __init__(self, manifold_module: ProperVelocity, num_features: int, **kwargs):
    validate_pv_manifold(manifold_module, required_methods=_GYRO_BN_METHODS)
    super().__init__(manifold_module, num_features, **kwargs)

hyperbolix.nn_layers.HyperboloidGyroRMSNorm

HyperboloidGyroRMSNorm(
    manifold_module: Hyperboloid,
    num_features: int,
    **kwargs,
)

Bases: GyroRMSNormBase

Gyro radial RMSNorm for the Hyperboloid (Lorentz) model.

Inputs are ambient (..., D+1) hyperboloid points; num_features = D.

Source code in hyperbolix/nn_layers/gyro_normalization.py

def __init__(self, manifold_module: Hyperboloid, num_features: int, **kwargs):
    validate_hyperboloid_manifold(manifold_module, required_methods=_GYRO_RMS_METHODS)
    super().__init__(manifold_module, num_features, **kwargs)

hyperbolix.nn_layers.ProperVelocityGyroRMSNorm

ProperVelocityGyroRMSNorm(
    manifold_module: ProperVelocity,
    num_features: int,
    **kwargs,
)

Bases: GyroRMSNormBase

Gyro radial RMSNorm for the Proper Velocity (PV) model.

Inputs are (..., D) PV points; num_features = D.

Source code in hyperbolix/nn_layers/gyro_normalization.py

def __init__(self, manifold_module: ProperVelocity, num_features: int, **kwargs):
    validate_pv_manifold(manifold_module, required_methods=_GYRO_RMS_METHODS)
    super().__init__(manifold_module, num_features, **kwargs)

hyperbolix.nn_layers.PoincareGyroRMSNorm

PoincareGyroRMSNorm(
    manifold_module: Poincare, num_features: int, **kwargs
)

Bases: GyroRMSNormBase

Gyro radial RMSNorm for the Poincaré ball model.

Inputs are on-ball points (..., D) (no time coordinate); num_features = D. The Poincaré ball has no time dimension, so embed_spatial_0 is the identity and the ambient/spatial distinction collapses. The radial rescale uses Möbius scalar_mul exactly as the Hyperboloid/PV variants use Lorentz/PV scalar_mul; all three scale the geodesic radius dist_0(x) -> |r|·dist_0(x) linearly, so a single r = gamma / (dist_0(x) + eps) sends every sample to radius ≈ gamma.

This is the per-sample, batch-independent counterpart to the (tangent-space, batch-statistics) :class:PoincareBatchNorm2D — valid at batch size 1 and identical in train and eval.

Source code in hyperbolix/nn_layers/gyro_normalization.py

def __init__(self, manifold_module: Poincare, num_features: int, **kwargs):
    validate_poincare_manifold(manifold_module, required_methods=_GYRO_RMS_METHODS)
    super().__init__(manifold_module, num_features, **kwargs)

HRC Normalization & Dropout (Hyperboloid)

HRC-wrapped normalizers and dropout for the Hyperboloid (Hypformer). They apply a Euclidean operation to the space components and reconstruct the time component, with curvature-change support. FGGMeanOnlyBatchNorm pairs with FGGLinear(use_weight_norm=True).

hyperbolix.nn_layers.HRCBatchNorm

HRCBatchNorm(
    num_features: int,
    *,
    rngs: Rngs,
    momentum: float = 0.99,
    epsilon: float = 1e-05,
    eps: float = 1e-07,
    param_dtype: DTypeLike = jnp.float32,
)

Bases: Module

Hyperbolic Regularization Component with batch normalization.

Applies batch normalization to spatial components of hyperboloid points, then reconstructs the time component for the output curvature.

Parameters:

Name	Type	Description	Default
num_features	int	Number of spatial features to normalize.	required
rngs	Rngs	Random number generators for parameter initialization.	required
momentum	float	Momentum for running statistics (default: 0.99).	0.99
epsilon	float	Small value for numerical stability in batch norm (default: 1e-5).	1e-05
eps	float	Small value for numerical stability in HRC (default: 1e-7).	1e-07
param_dtype	DTypeLike	Storage dtype of the trainable parameters (default: jnp.float32).	float32

Attributes:

Name	Type	Description
bn	BatchNorm	Flax batch normalization.
eps	float	Numerical stability parameter for HRC.

Notes

Training vs Evaluation Mode: During training (use_running_average=False), batch norm computes statistics from the current batch and updates running averages. During evaluation (use_running_average=True), it uses the accumulated running statistics.

Examples:

>>> from flax import nnx
>>> from hyperbolix.nn_layers import HRCBatchNorm
>>>
>>> # Create batch norm layer
>>> bn = HRCBatchNorm(num_features=64, rngs=nnx.Rngs(0))
>>>
>>> # Training mode
>>> y_train = bn(x, c_in=1.0, c_out=2.0, use_running_average=False)
>>>
>>> # Evaluation mode
>>> y_eval = bn(x, c_in=1.0, c_out=2.0, use_running_average=True)

Source code in hyperbolix/nn_layers/hyperboloid_regularization.py

def __init__(
    self,
    num_features: int,
    *,
    rngs: nnx.Rngs,
    momentum: float = 0.99,
    epsilon: float = 1e-5,
    eps: float = 1e-7,
    param_dtype: DTypeLike = jnp.float32,
):
    self.bn = nnx.BatchNorm(
        num_features,
        momentum=momentum,
        epsilon=epsilon,
        param_dtype=param_dtype,
        rngs=rngs,
    )
    self.eps = eps

__call__

__call__(
    x: Float[Array, "batch dim_plus_1"],
    c_in: float = 1.0,
    c_out: float = 1.0,
    use_running_average: bool | None = None,
) -> Float[Array, "batch dim_plus_1"]

Apply HRC batch normalization.

Parameters:

Name	Type	Description	Default
x	Array of shape (batch, dim+1)	Input points on hyperboloid with curvature c_in.	required
c_in	float	Input curvature (default: 1.0).	1.0
c_out	float	Output curvature (default: 1.0).	1.0
use_running_average	bool or None	If True, use running statistics (eval mode). If False, use batch statistics (train mode). If None, use the default set during initialization.	None

Returns:

Name	Type	Description
y	Array of shape (batch, dim+1)	Output points on hyperboloid with curvature c_out.

Source code in hyperbolix/nn_layers/hyperboloid_regularization.py

def __call__(
    self,
    x: Float[Array, "batch dim_plus_1"],
    c_in: float = 1.0,
    c_out: float = 1.0,
    use_running_average: bool | None = None,
) -> Float[Array, "batch dim_plus_1"]:
    """Apply HRC batch normalization.

    Parameters
    ----------
    x : Array of shape (batch, dim+1)
        Input points on hyperboloid with curvature c_in.
    c_in : float, optional
        Input curvature (default: 1.0).
    c_out : float, optional
        Output curvature (default: 1.0).
    use_running_average : bool or None, optional
        If True, use running statistics (eval mode).
        If False, use batch statistics (train mode).
        If None, use the default set during initialization.

    Returns
    -------
    y : Array of shape (batch, dim+1)
        Output points on hyperboloid with curvature c_out.
    """

    def bn_fn(z):
        return self.bn(z, use_running_average=use_running_average)

    return hrc(x, bn_fn, c_in, c_out, self.eps)

hyperbolix.nn_layers.HRCLayerNorm

HRCLayerNorm(
    num_features: int,
    *,
    rngs: Rngs,
    epsilon: float = 1e-05,
    eps: float = 1e-07,
    param_dtype: DTypeLike = jnp.float32,
)

Bases: Module

Hyperbolic Regularization Component with layer normalization.

Applies layer normalization to spatial components of hyperboloid points, then reconstructs the time component for the output curvature.

Parameters:

Name	Type	Description	Default
num_features	int	Number of spatial features to normalize.	required
rngs	Rngs	Random number generators for parameter initialization.	required
epsilon	float	Small value for numerical stability in layer norm (default: 1e-5).	1e-05
eps	float	Small value for numerical stability in HRC (default: 1e-7).	1e-07
param_dtype	DTypeLike	Storage dtype of the trainable parameters (default: jnp.float32).	float32

Attributes:

Name	Type	Description
ln	LayerNorm	Flax layer normalization.
eps	float	Numerical stability parameter for HRC.

Examples:

>>> from flax import nnx
>>> from hyperbolix.nn_layers import HRCLayerNorm
>>>
>>> ln = HRCLayerNorm(num_features=64, rngs=nnx.Rngs(0))
>>> y = ln(x, c_in=1.0, c_out=2.0)

Source code in hyperbolix/nn_layers/hyperboloid_regularization.py

def __init__(
    self,
    num_features: int,
    *,
    rngs: nnx.Rngs,
    epsilon: float = 1e-5,
    eps: float = 1e-7,
    param_dtype: DTypeLike = jnp.float32,
):
    self.ln = nnx.LayerNorm(num_features, epsilon=epsilon, param_dtype=param_dtype, rngs=rngs)
    self.eps = eps

__call__

__call__(
    x: Float[Array, "batch dim_plus_1"],
    c_in: float = 1.0,
    c_out: float = 1.0,
) -> Float[Array, "batch dim_plus_1"]

Apply HRC layer normalization.

Parameters:

Name	Type	Description	Default
x	Array of shape (batch, dim+1)	Input points on hyperboloid with curvature c_in.	required
c_in	float	Input curvature (default: 1.0).	1.0
c_out	float	Output curvature (default: 1.0).	1.0

Returns:

Name	Type	Description
y	Array of shape (batch, dim+1)	Output points on hyperboloid with curvature c_out.

Source code in hyperbolix/nn_layers/hyperboloid_regularization.py

def __call__(
    self,
    x: Float[Array, "batch dim_plus_1"],
    c_in: float = 1.0,
    c_out: float = 1.0,
) -> Float[Array, "batch dim_plus_1"]:
    """Apply HRC layer normalization.

    Parameters
    ----------
    x : Array of shape (batch, dim+1)
        Input points on hyperboloid with curvature c_in.
    c_in : float, optional
        Input curvature (default: 1.0).
    c_out : float, optional
        Output curvature (default: 1.0).

    Returns
    -------
    y : Array of shape (batch, dim+1)
        Output points on hyperboloid with curvature c_out.
    """
    return hrc(x, self.ln, c_in, c_out, self.eps)

hyperbolix.nn_layers.HRCRMSNorm

HRCRMSNorm(
    num_features: int,
    *,
    rngs: Rngs,
    epsilon: float = 1e-06,
    eps: float = 1e-07,
    param_dtype: DTypeLike = jnp.float32,
)

Bases: Module

Hyperbolic Regularization Component with RMS normalization.

Applies RMS normalization to spatial components of hyperboloid points, then reconstructs the time component for the output curvature. RMSNorm is a simpler and faster variant of LayerNorm that only normalizes by the root mean square without centering.

Parameters:

Name	Type	Description	Default
num_features	int	Number of spatial features to normalize.	required
rngs	Rngs	Random number generators for parameter initialization.	required
epsilon	float	Small value for numerical stability in RMS norm (default: 1e-6).	1e-06
eps	float	Small value for numerical stability in HRC (default: 1e-7).	1e-07
param_dtype	DTypeLike	Storage dtype of the trainable parameters (default: jnp.float32).	float32

Attributes:

Name	Type	Description
rms	RMSNorm	Flax RMS normalization.
eps	float	Numerical stability parameter for HRC.

Notes

RMSNorm Formula: y = x / RMS(x) * scale, where RMS(x) = sqrt(mean(x^2) + epsilon)

RMSNorm is more efficient than LayerNorm as it skips mean subtraction and is commonly used in modern transformers (e.g., LLaMA, GPT-NeoX).

Examples:

>>> from flax import nnx
>>> from hyperbolix.nn_layers import HRCRMSNorm
>>>
>>> rms = HRCRMSNorm(num_features=64, rngs=nnx.Rngs(0))
>>> y = rms(x, c_in=1.0, c_out=2.0)

Source code in hyperbolix/nn_layers/hyperboloid_regularization.py

def __init__(
    self,
    num_features: int,
    *,
    rngs: nnx.Rngs,
    epsilon: float = 1e-6,
    eps: float = 1e-7,
    param_dtype: DTypeLike = jnp.float32,
):
    self.rms = nnx.RMSNorm(num_features, epsilon=epsilon, param_dtype=param_dtype, rngs=rngs)
    self.eps = eps

__call__

__call__(
    x: Float[Array, "batch dim_plus_1"],
    c_in: float = 1.0,
    c_out: float = 1.0,
) -> Float[Array, "batch dim_plus_1"]

Apply HRC RMS normalization.

Parameters:

Name	Type	Description	Default
x	Array of shape (batch, dim+1)	Input points on hyperboloid with curvature c_in.	required
c_in	float	Input curvature (default: 1.0).	1.0
c_out	float	Output curvature (default: 1.0).	1.0

Returns:

Name	Type	Description
y	Array of shape (batch, dim+1)	Output points on hyperboloid with curvature c_out.

Source code in hyperbolix/nn_layers/hyperboloid_regularization.py

def __call__(
    self,
    x: Float[Array, "batch dim_plus_1"],
    c_in: float = 1.0,
    c_out: float = 1.0,
) -> Float[Array, "batch dim_plus_1"]:
    """Apply HRC RMS normalization.

    Parameters
    ----------
    x : Array of shape (batch, dim+1)
        Input points on hyperboloid with curvature c_in.
    c_in : float, optional
        Input curvature (default: 1.0).
    c_out : float, optional
        Output curvature (default: 1.0).

    Returns
    -------
    y : Array of shape (batch, dim+1)
        Output points on hyperboloid with curvature c_out.
    """
    return hrc(x, self.rms, c_in, c_out, self.eps)

hyperbolix.nn_layers.HRCDropout

HRCDropout(rate: float, *, rngs: Rngs, eps: float = 1e-07)

Bases: Module

Hyperbolic Regularization Component with dropout.

Applies dropout to spatial components of hyperboloid points, then reconstructs the time component for the output curvature.

Parameters:

Name	Type	Description	Default
rate	float	Dropout probability (fraction of units to drop).	required
rngs	Rngs	Random number generators for dropout.	required
eps	float	Small value for numerical stability (default: 1e-7).	1e-07

Attributes:

Name	Type	Description
dropout	Dropout	Flax dropout layer.
eps	float	Numerical stability parameter.

Examples:

>>> from flax import nnx
>>> from hyperbolix.nn_layers import HRCDropout
>>>
>>> dropout = HRCDropout(rate=0.1, rngs=nnx.Rngs(dropout=42))
>>> y = dropout(x, c_in=1.0, c_out=1.0, deterministic=False)

Source code in hyperbolix/nn_layers/hyperboloid_regularization.py

def __init__(self, rate: float, *, rngs: nnx.Rngs, eps: float = 1e-7):
    self.dropout = nnx.Dropout(rate, rngs=rngs)
    self.eps = eps

__call__

__call__(
    x: Float[Array, "batch dim_plus_1"],
    c_in: float = 1.0,
    c_out: float = 1.0,
    deterministic: bool = False,
) -> Float[Array, "batch dim_plus_1"]

Apply HRC dropout.

Parameters:

Name	Type	Description	Default
x	Array of shape (batch, dim+1)	Input points on hyperboloid with curvature c_in.	required
c_in	float	Input curvature (default: 1.0).	1.0
c_out	float	Output curvature (default: 1.0).	1.0
deterministic	bool	If True, no dropout is applied (for evaluation mode).	False

Returns:

Name	Type	Description
y	Array of shape (batch, dim+1)	Output points on hyperboloid with curvature c_out.

Source code in hyperbolix/nn_layers/hyperboloid_regularization.py

def __call__(
    self,
    x: Float[Array, "batch dim_plus_1"],
    c_in: float = 1.0,
    c_out: float = 1.0,
    deterministic: bool = False,
) -> Float[Array, "batch dim_plus_1"]:
    """Apply HRC dropout.

    Parameters
    ----------
    x : Array of shape (batch, dim+1)
        Input points on hyperboloid with curvature c_in.
    c_in : float, optional
        Input curvature (default: 1.0).
    c_out : float, optional
        Output curvature (default: 1.0).
    deterministic : bool, optional
        If True, no dropout is applied (for evaluation mode).

    Returns
    -------
    y : Array of shape (batch, dim+1)
        Output points on hyperboloid with curvature c_out.
    """

    def drop_fn(z):
        return self.dropout(z, deterministic=deterministic)

    return hrc(x, drop_fn, c_in, c_out, self.eps)

hyperbolix.nn_layers.FGGMeanOnlyBatchNorm

FGGMeanOnlyBatchNorm(
    num_features: int,
    *,
    momentum: float = 0.99,
    eps: float = 1e-07,
    param_dtype: DTypeLike = jnp.float32,
)

Bases: Module

Mean-only batch normalization for FGG-LNN layers.

Subtracts the batch mean from spatial components without dividing by variance, then adds a learnable bias. This avoids the instability that standard BatchNorm causes in hyperbolic space (exponentially large/small variance denominators) while still centering activations for stable training.

Designed to pair with weight normalization (FGGLinear(use_weight_norm=True)), which handles magnitude control. Together they replace standard BatchNorm in fully hyperbolic FGG-LNN networks.

Formula (applied to spatial components via HRC): y = x - E[x] + bias

Parameters:

Name	Type	Description	Default
num_features	int	Number of spatial features (D, not D+1).	required
momentum	float	Exponential moving average momentum for running mean (default: 0.99). Update: running_mean = momentum * running_mean + (1 - momentum) * batch_mean.	0.99
eps	float	Numerical stability floor for HRC time reconstruction (default: 1e-7).	1e-07
param_dtype	DTypeLike	Storage dtype of the learnable bias and running mean (default: jnp.float32).	float32

References

Klis et al. "Fast and Geometrically Grounded Lorentz Neural Networks" (2026), §4.4. Salimans & Kingma "Weight Normalization" (2016).

Source code in hyperbolix/nn_layers/hyperboloid_regularization.py

def __init__(
    self,
    num_features: int,
    *,
    momentum: float = 0.99,
    eps: float = 1e-7,
    param_dtype: DTypeLike = jnp.float32,
):
    self.num_features = num_features
    self.momentum = momentum
    self.eps = eps

    # Pin the bias/running-mean storage dtype. NNX casts assignments back
    # to a Variable's init dtype, so a float64 init (which a bare
    # jnp.zeros becomes under global jax_enable_x64) would persist float64
    # state regardless of the working precision. Defaulting to float32
    # matches the library-wide param_dtype convention; pass
    # param_dtype=jnp.float64 for a fully float64 network.
    # Learnable bias (shift after centering)
    self.bias = nnx.Param(jnp.zeros((num_features,), dtype=param_dtype))
    # Running mean for eval mode
    self.running_mean = nnx.BatchStat(jnp.zeros((num_features,), dtype=param_dtype))

__call__

__call__(
    x: Float[Array, "batch dim_plus_1"],
    c_in: float = 1.0,
    c_out: float = 1.0,
    use_running_average: bool = False,
) -> Float[Array, "batch dim_plus_1"]

Apply mean-only batch normalization via HRC.

Parameters:

Name	Type	Description	Default
x	(Array, shape(B, D + 1) or (B, ..., D + 1))	Input points on hyperboloid with curvature c_in.	required
c_in	float	Input curvature (default: 1.0).	1.0
c_out	float	Output curvature (default: 1.0).	1.0
use_running_average	bool	If True, use running mean (eval mode). If False, compute from batch and update running mean (train mode). Default: False.	False

Returns:

Type	Description
Array, same shape as input	Points on hyperboloid with curvature c_out.

Source code in hyperbolix/nn_layers/hyperboloid_regularization.py

def __call__(
    self,
    x: Float[Array, "batch dim_plus_1"],
    c_in: float = 1.0,
    c_out: float = 1.0,
    use_running_average: bool = False,
) -> Float[Array, "batch dim_plus_1"]:
    """Apply mean-only batch normalization via HRC.

    Parameters
    ----------
    x : Array, shape (B, D+1) or (B, ..., D+1)
        Input points on hyperboloid with curvature ``c_in``.
    c_in : float, optional
        Input curvature (default: 1.0).
    c_out : float, optional
        Output curvature (default: 1.0).
    use_running_average : bool, optional
        If True, use running mean (eval mode). If False, compute from
        batch and update running mean (train mode). Default: False.

    Returns
    -------
    Array, same shape as input
        Points on hyperboloid with curvature ``c_out``.
    """

    def mean_only_bn(z):
        # z: spatial components, shape (..., D). Cast state to the working
        # dtype so float64 bias/running-mean (from global jax_enable_x64)
        # don't promote the float32 computation or persist as float64 state.
        bias = self.bias[...].astype(z.dtype)
        if use_running_average:
            mean = self.running_mean[...].astype(z.dtype)
        else:
            # Compute mean over all dims except the last (feature) dim
            # Flatten leading dims for mean computation
            z_flat = z.reshape(-1, z.shape[-1])  # (N, D)
            mean = jnp.mean(z_flat, axis=0)  # (D,)

            # Update running mean (EMA). Compute in the working dtype, then
            # cast explicitly to the stat's own dtype: NNX casts assignments
            # back to the Variable's init dtype anyway, and doing it
            # explicitly avoids an unsafe-scatter FutureWarning when a float64
            # network feeds the default float32 stat.
            rm = self.running_mean[...]
            new_rm = self.momentum * rm.astype(z.dtype) + (1.0 - self.momentum) * mean
            self.running_mean[...] = new_rm.astype(rm.dtype)

        return z - mean + bias

    return hrc(x, mean_only_bn, c_in, c_out, self.eps)

Euclidean input scaling (hybrid networks)

HyperPPFeatureScaling is not an on-manifold normalizer — it runs entirely in Euclidean space, applied after the last Euclidean layer and before expmap_0 to either manifold. It is RMSNorm-based (RMSNorm + Lipschitz activation + dimension scaling + optional learned rescaling), bounding the output norm via rho_max = atanh(alpha) / sqrt(c). See the boundary-pattern recipe.

hyperbolix.nn_layers.HyperPPFeatureScaling

HyperPPFeatureScaling(
    dim: int,
    *,
    alpha: float | None = None,
    activation: Callable | None = jax.nn.tanh,
    param_dtype: DTypeLike = jnp.float32,
    rngs: Rngs,
)

Bases: Module

Hyper++ feature scaling for hybrid Euclidean-hyperbolic networks.

Parameters:

Name	Type	Description	Default
dim	int	Embedding dimension (needed for nnx.Linear and 1/sqrt(d) scaling).	required
alpha	float or None	Value in (0, 1) enabling learned rescaling. None disables it and the layer is entirely parameter-free. Default: None.	None
activation	Callable or None	Lipschitz activation function. Default: jax.nn.tanh. None to skip.	tanh
param_dtype	DTypeLike	Storage dtype of the trainable parameters (default: jnp.float32).	float32
rngs	Rngs	Random number generators for sub-layers.	required

Examples:

>>> from flax import nnx
>>> from hyperbolix.nn_layers import HyperPPFeatureScaling
>>>
>>> # Parameter-free mode
>>> layer = HyperPPFeatureScaling(dim=64, rngs=nnx.Rngs(0))
>>> y = layer(x, c=1.0)
>>>
>>> # Learned rescaling mode
>>> layer = HyperPPFeatureScaling(dim=64, alpha=0.9, rngs=nnx.Rngs(0))
>>> y = layer(x, c=0.1)

Source code in hyperbolix/nn_layers/hybrid_regularization.py

def __init__(
    self,
    dim: int,
    *,
    alpha: float | None = None,
    activation: Callable | None = jax.nn.tanh,
    param_dtype: DTypeLike = jnp.float32,
    rngs: nnx.Rngs,
):
    if alpha is not None:
        if not (0.0 < alpha < 1.0):
            msg = f"alpha must be in (0, 1), got {alpha}"
            raise ValueError(msg)
        self._atanh_alpha = math.atanh(alpha)
        self.xi_theta = nnx.Linear(dim, 1, param_dtype=param_dtype, rngs=rngs)
    else:
        self._atanh_alpha = None
        self.xi_theta = None

    self.rms_norm = nnx.RMSNorm(dim, use_scale=False, param_dtype=param_dtype, rngs=rngs)
    self._inv_sqrt_dim = 1.0 / math.sqrt(dim)
    self.activation = activation

__call__

__call__(
    x_BD: Float[Array, "B D"], c: float = 1.0
) -> Float[Array, "B D"]

Apply Hyper++ feature scaling pipeline.

Parameters:

Name	Type	Description	Default
x_BD	Array of shape (B, D)	Euclidean features from the last Euclidean layer.	required
c	float	Curvature parameter (positive). Passed at call time per library convention to support learnable curvature.	1.0

Returns:

Type	Description
Array of shape (B, D)	Scaled features ready for expmap_0.

Source code in hyperbolix/nn_layers/hybrid_regularization.py

def __call__(
    self,
    x_BD: Float[Array, "B D"],
    c: float = 1.0,
) -> Float[Array, "B D"]:
    """Apply Hyper++ feature scaling pipeline.

    Parameters
    ----------
    x_BD : Array of shape (B, D)
        Euclidean features from the last Euclidean layer.
    c : float
        Curvature parameter (positive). Passed at call time per
        library convention to support learnable curvature.

    Returns
    -------
    Array of shape (B, D)
        Scaled features ready for expmap_0.
    """
    # Step 1: RMSNorm (parameter-free, no learned scale)
    x_BD = self.rms_norm(x_BD)

    # Step 2: Lipschitz activation
    if self.activation is not None:
        x_BD = self.activation(x_BD)

    # Step 3: Dimension scaling
    x_BD = x_BD * self._inv_sqrt_dim

    # Step 4: Learned rescaling
    if self._atanh_alpha is not None:
        assert self.xi_theta is not None
        rho_max = self._atanh_alpha / jnp.sqrt(c)  # scalar
        scale_B1 = rho_max * jax.nn.sigmoid(self.xi_theta(x_BD))  # (B, 1)
        x_BD = scale_B1 * x_BD

    return x_BD

Example

import jax
from hyperbolix.manifolds import Hyperboloid, ProperVelocity
from hyperbolix.nn_layers import HyperboloidGyroBatchNorm, HyperboloidGyroRMSNorm, ProperVelocityGyroRMSNorm

c = 1.0
hyp = Hyperboloid()
bn = HyperboloidGyroBatchNorm(hyp, num_features=16)   # num_features = spatial D
rms = HyperboloidGyroRMSNorm(hyp, num_features=16)
spatial = jax.random.normal(jax.random.PRNGKey(0), (8, 16)) * 0.2
x_amb = jax.vmap(hyp.expmap_0, in_axes=(0, None))(jax.vmap(hyp.embed_spatial_0)(spatial), c)  # (8, 17)
h = bn(x_amb, c=c)                            # training (running stats updated)
h = bn(x_amb, c=c, use_running_average=True)  # evaluation (frozen stats)
h = rms(x_amb, c=c)                           # per-sample, batch-free

pv_rms = ProperVelocityGyroRMSNorm(ProperVelocity(), num_features=16, use_bias=True)
h_pv = pv_rms(jax.random.normal(jax.random.PRNGKey(1), (8, 16)) * 0.2, c=c)  # (8, 16)