Linear Layers

Fully-connected hyperbolic layers across the Poincaré, Hyperboloid, and Proper Velocity models. For which linear layer to pick (defaults, speed/expressiveness trade-offs) see the NN Layers guide.

Channel convention

Hyperboloid linear layers (HTCLinear, FGGLinear, HypLinearHyperboloid*) take ambient in_features = d+1. Poincaré and PV layers take spatial in_dim = d. HTCLinear is the exception whose out_features is spatial — its output shape is (B, out_features + 1).

Poincaré

hyperbolix.nn_layers.HypLinearPoincare

HypLinearPoincare(
    manifold_module: Manifold,
    in_dim: int,
    out_dim: int,
    *,
    rngs: Rngs,
    input_space: str = "manifold",
    curvature: float | Callable[[], Any] = 1.0,
    param_dtype: DTypeLike = jnp.float32,
)

Bases: Module

Hyperbolic Neural Networks fully connected layer (Poincaré ball model).

Computation steps: 0) Project the input tensor to the tangent space (optional) 1) Perform matrix vector multiplication in the tangent space at the origin. 2) Map the result to the manifold. 3) Add the manifold bias to the result.

Parameters:

Name	Type	Description	Default
manifold_module	object	Class-based Poincare manifold instance	required
in_dim	int	Dimension of the input space	required
out_dim	int	Dimension of the output space	required
rngs	Rngs	Random number generators for parameter initialization	required
input_space	str	Type of the input tensor, either 'tangent' or 'manifold' (default: 'manifold'). Note: This is a static configuration - changing it after initialization requires recompilation.	'manifold'
curvature	float or callable	Curvature tag for the manifold-valued bias parameter (default: 1.0). The Riemannian optimizer uses this value for the bias update (egrad2rgrad, expmap, parallel transport), so it MUST match the c passed at __call__ time — a mismatch silently applies the wrong Riemannian correction. For learnable curvature, pass a callable returning the current value (e.g. lambda: model.curvature()).	1.0
param_dtype	DTypeLike	Storage dtype of the trainable parameters (default: jnp.float32). Compute precision of manifold operations is set by manifold.dtype.	float32

Notes

JIT Compatibility: This layer is designed to work with nnx.jit. Configuration parameters (input_space) are treated as static and will be baked into the compiled function. Changing these values after JIT compilation will trigger automatic recompilation.

References

Ganea Octavian, Gary Bécigneul, and Thomas Hofmann. "Hyperbolic neural networks." Advances in neural information processing systems 31 (2018).

Source code in hyperbolix/nn_layers/poincare_linear.py

def __init__(
    self,
    manifold_module: Manifold,
    in_dim: int,
    out_dim: int,
    *,
    rngs: nnx.Rngs,
    input_space: str = "manifold",
    curvature: float | Callable[[], Any] = 1.0,
    param_dtype: DTypeLike = jnp.float32,
):
    if input_space not in ["tangent", "manifold"]:
        raise ValueError(f"input_space must be either 'tangent' or 'manifold', got '{input_space}'")

    # Static configuration (treated as compile-time constants for JIT)
    validate_poincare_manifold(
        manifold_module,
        required_methods=("proj", "addition", "expmap_0", "logmap_0", "compute_mlr_pp"),
    )
    self.manifold = manifold_module
    self.in_dim = in_dim
    self.out_dim = out_dim
    self.input_space = input_space

    # Trainable parameters
    # Tangent space weight (Euclidean) - initialized with std = 1/sqrt(fan_in)
    # to prevent outputs from saturating at the Poincaré ball boundary
    std = 1.0 / jnp.sqrt(in_dim)
    self.kernel = nnx.Param(jax.random.normal(rngs.params(), (out_dim, in_dim), dtype=param_dtype) * std)
    # Manifold bias (initialized to small random values to avoid gradient issues at origin)
    self.bias = ManifoldParam(
        jax.random.normal(rngs.params(), (out_dim,), dtype=param_dtype) * 0.01,
        manifold=self.manifold,
        curvature=curvature,
    )

__call__

__call__(
    x: Float[Array, "batch in_dim"], c: float = 1.0
) -> Float[Array, "batch out_dim"]

Forward pass through the hyperbolic linear layer.

Parameters:

Name	Type	Description	Default
x	Array of shape (batch, in_dim)	Input tensor where the hyperbolic_axis is last	required
c	float	Manifold curvature (default: 1.0)	1.0

Returns:

Name	Type	Description
res	Array of shape (batch, out_dim)	Output on the Poincaré ball manifold

Source code in hyperbolix/nn_layers/poincare_linear.py

def __call__(
    self,
    x: Float[Array, "batch in_dim"],
    c: float = 1.0,
) -> Float[Array, "batch out_dim"]:
    """
    Forward pass through the hyperbolic linear layer.

    Parameters
    ----------
    x : Array of shape (batch, in_dim)
        Input tensor where the hyperbolic_axis is last
    c : float
        Manifold curvature (default: 1.0)

    Returns
    -------
    res : Array of shape (batch, out_dim)
        Output on the Poincaré ball manifold
    """
    # Project bias to manifold
    bias_O = self.manifold.proj(self.bias[...], c)

    # Map to tangent space if needed (static branch - JIT friendly)
    if self.input_space == "manifold":
        x_BI = jax.vmap(self.manifold.logmap_0, in_axes=(0, None), out_axes=0)(x, c)
    else:
        x_BI = x

    # Matrix-vector multiplication in tangent space at origin. Cast the
    # kernel to the input dtype so float64 weights (from global
    # jax_enable_x64) don't promote a float32 computation to float64.
    kernel_OI = self.kernel[...].astype(x_BI.dtype)
    x_BO = jnp.einsum("bi,oi->bo", x_BI, kernel_OI)  # (B, I) @ (I, O) -> (B, O)

    # Map back to manifold
    x_BO = jax.vmap(self.manifold.expmap_0, in_axes=(0, None), out_axes=0)(x_BO, c)

    # Manifold bias addition (Möbius addition for Poincaré)
    res_BO = jax.vmap(self.manifold.addition, in_axes=(0, None, None), out_axes=0)(x_BO, bias_O, c)
    return res_BO

hyperbolix.nn_layers.HypLinearPoincarePP

HypLinearPoincarePP(
    manifold_module: Poincare,
    in_dim: int,
    out_dim: int,
    *,
    rngs: Rngs,
    input_space: str = "manifold",
    clamping_factor: float = 1.0,
    smoothing_factor: float = 50.0,
    param_dtype: DTypeLike = jnp.float32,
)

Bases: Module

Hyperbolic Neural Networks ++ fully connected layer (Poincaré ball model).

Computation steps: 0) Project the input tensor onto the manifold (optional) 1) Compute the multinomial linear regression score(s) 2) Calculate the generalized linear transformation from the regression score(s)

Parameters:

Name	Type	Description	Default
manifold_module	object	Class-based Poincare manifold instance	required
in_dim	int	Dimension of the input space	required
out_dim	int	Dimension of the output space	required
rngs	Rngs	Random number generators for parameter initialization	required
input_space	str	Type of the input tensor, either 'tangent' or 'manifold' (default: 'manifold'). Note: This is a static configuration - changing it after initialization requires recompilation.	'manifold'
clamping_factor	float	Clamping factor for the multinomial linear regression output (default: 1.0)	1.0
smoothing_factor	float	Smoothing factor for the multinomial linear regression output (default: 50.0)	50.0
param_dtype	DTypeLike	Storage dtype of the trainable parameters (default: jnp.float32). Compute precision of manifold operations is set by manifold.dtype.	float32

Notes

JIT Compatibility: This layer is designed to work with nnx.jit. Configuration parameters (input_space, clamping_factor, smoothing_factor) are treated as static and will be baked into the compiled function.

References

Shimizu Ryohei, Yusuke Mukuta, and Tatsuya Harada. "Hyperbolic neural networks++." arXiv preprint arXiv:2006.08210 (2020).

Source code in hyperbolix/nn_layers/poincare_linear.py

def __init__(
    self,
    manifold_module: Poincare,
    in_dim: int,
    out_dim: int,
    *,
    rngs: nnx.Rngs,
    input_space: str = "manifold",
    clamping_factor: float = 1.0,
    smoothing_factor: float = 50.0,
    param_dtype: DTypeLike = jnp.float32,
):
    if input_space not in ["tangent", "manifold"]:
        raise ValueError(f"input_space must be either 'tangent' or 'manifold', got '{input_space}'")

    # Static configuration (treated as compile-time constants for JIT)
    validate_poincare_manifold(
        manifold_module,
        required_methods=("proj", "addition", "expmap_0", "logmap_0", "compute_mlr_pp"),
    )
    self.manifold = manifold_module
    self.in_dim = in_dim
    self.out_dim = out_dim
    self.input_space = input_space
    self.clamping_factor = clamping_factor
    self.smoothing_factor = smoothing_factor

    # Trainable parameters
    # Tangent space weight - initialized with std = 1/sqrt(fan_in)
    # to prevent outputs from saturating at the Poincaré ball boundary
    std = 1.0 / jnp.sqrt(in_dim)
    self.kernel = nnx.Param(jax.random.normal(rngs.params(), (out_dim, in_dim), dtype=param_dtype) * std)
    # Scalar bias
    self.bias = nnx.Param(jnp.zeros((out_dim, 1), dtype=param_dtype))

__call__

__call__(
    x: Float[Array, "batch in_dim"], c: float = 1.0
) -> Float[Array, "batch out_dim"]

Forward pass through the HNN++ hyperbolic linear layer.

Parameters:

Name	Type	Description	Default
x	Array of shape (batch, in_dim)	Input tensor where the hyperbolic_axis is last	required
c	float	Manifold curvature (default: 1.0)	1.0

Returns:

Name	Type	Description
res	Array of shape (batch, out_dim)	Output on the Poincaré ball manifold

Source code in hyperbolix/nn_layers/poincare_linear.py

def __call__(
    self,
    x: Float[Array, "batch in_dim"],
    c: float = 1.0,
) -> Float[Array, "batch out_dim"]:
    """
    Forward pass through the HNN++ hyperbolic linear layer.

    Parameters
    ----------
    x : Array of shape (batch, in_dim)
        Input tensor where the hyperbolic_axis is last
    c : float
        Manifold curvature (default: 1.0)

    Returns
    -------
    res : Array of shape (batch, out_dim)
        Output on the Poincaré ball manifold
    """
    return _poincare_pp_forward(
        x,
        self.kernel[...],
        self.bias[...],
        self.manifold,
        c,
        self.input_space,
        self.clamping_factor,
        self.smoothing_factor,
    )

hyperbolix.nn_layers.HypLinearPoincareBusemann

HypLinearPoincareBusemann(
    manifold_module: Poincare,
    in_dim: int,
    out_dim: int,
    *,
    rngs: Rngs,
    input_space: str = "manifold",
    activation: Callable[[Array], Array] | None = None,
    v_max: float = 10.0,
    use_gyro_bias: bool = False,
    param_dtype: DTypeLike = jnp.float32,
)

Bases: Module

Busemann fully connected (BFC) layer (Poincaré ball model).

Maps a Poincaré ball point to a Poincaré ball point (Chen et al. 2026, Thm. 4.1). Each output coordinate is the activated Busemann logit φ(-alpha_k·B^{v_k}(x) + b_k); the point is recovered via ω = sinh(√c·φ(u))/√c, y = ω / (1 + √(1 + c‖ω‖²)), then projected.

Parameters:

Name	Type	Description	Default
manifold_module	Poincare	Class-based Poincaré manifold instance.	required
in_dim	int	Input spatial dimension (the ball has no time component).	required
out_dim	int	Output spatial dimension.	required
rngs	Rngs	Random number generators for parameter initialization.	required
input_space	str	"manifold" (default) or "tangent" (lift via expmap_0 first). Static for JIT.	'manifold'
activation	Callable or None	Optional Euclidean activation φ applied to the Busemann logits (default: identity). Avoid relu when stacking several of these layers on high-dimensional input — same origin-collapse risk as :class:HypLinearHyperboloidBusemann (see its activation doc); tanh avoids it.	None
v_max	float	Output-side guard: √c·φ(u) is hard-clipped to ±v_max before the sinh map (default: 10.0).	10.0
use_gyro_bias	bool	If True, add a learnable intrinsic bias y ← y ⊕ exp_0(b) via Möbius addition, b zero-initialized (gyrogroup identity ⇒ no-op at init) (default: False).	False
param_dtype	DTypeLike	Storage dtype of the trainable parameters (default: jnp.float32).	float32

References

Chen, Schölkopf, and Sebe. "Hyperbolic Busemann Neural Networks." 2026, Thm. 4.1.

Source code in hyperbolix/nn_layers/busemann_linear.py

def __init__(
    self,
    manifold_module: Poincare,
    in_dim: int,
    out_dim: int,
    *,
    rngs: nnx.Rngs,
    input_space: str = "manifold",
    activation: Callable[[Array], Array] | None = None,
    v_max: float = 10.0,
    use_gyro_bias: bool = False,
    param_dtype: DTypeLike = jnp.float32,
):
    if input_space not in ["tangent", "manifold"]:
        raise ValueError(f"input_space must be either 'tangent' or 'manifold', got '{input_space}'")

    required_methods = ("expmap_0", "busemann")
    if use_gyro_bias:
        required_methods = ("expmap_0", "busemann", "addition")
    validate_poincare_manifold(manifold_module, required_methods=required_methods)
    _assert_v_max_safe(v_max)

    self.manifold = manifold_module
    self.in_dim = in_dim
    self.out_dim = out_dim
    self.input_space = input_space
    self.activation = activation
    self.v_max = v_max

    # No time component: spatial dim equals the model dim.
    in_spatial = in_dim
    out_spatial = out_dim
    self.kernel, self.log_scale, self.bias = init_weight_norm_params(
        rngs, out_spatial, in_spatial, std=(2 * in_spatial * out_spatial) ** -0.5, param_dtype=param_dtype
    )

    if use_gyro_bias:
        self.gyro_bias = nnx.Param(jnp.zeros((out_spatial,), dtype=param_dtype))
    else:
        self.gyro_bias = None

__call__

__call__(
    x: Float[Array, "batch in_dim"], c: float = 1.0
) -> Float[Array, "batch out_dim"]

Forward pass returning a Poincaré ball point, shape (batch, out_dim).

Source code in hyperbolix/nn_layers/busemann_linear.py

def __call__(
    self,
    x: Float[Array, "batch in_dim"],
    c: float = 1.0,
) -> Float[Array, "batch out_dim"]:
    """Forward pass returning a Poincaré ball point, shape (batch, out_dim)."""
    u_BO = busemann_score(self.manifold, x, self.kernel[...], self.log_scale[...], self.bias[...], c, self.input_space)
    if self.activation is not None:
        u_BO = self.activation(u_BO)
    y_BO = busemann_fc_poincare_output(u_BO, c, self.v_max)
    # Closed form already lands strictly inside the ball; proj guards against float rounding.
    y_BO = jax.vmap(self.manifold.proj, in_axes=(0, None))(y_BO, c)

    if self.gyro_bias is not None:
        # Paper notation F(x) ⊕ b (Sec. 4.2): Möbius addition with the bias point on the right.
        bias_point_O = self.manifold.expmap_0(self.gyro_bias[...].astype(y_BO.dtype), c)
        y_BO = jax.vmap(self.manifold.addition, in_axes=(0, None, None))(y_BO, bias_point_O, c)

    return y_BO

Hyperboloid

HypLinearHyperboloidPLFC is the point-to-hyperplane Lorentz FC (Shi et al. 2026), the Lorentz analog of the HNN++ formulation; the Busemann FC decides with point-to-horosphere distances. HTCLinear (the robust default) is documented here; the foundational htc() function lives in Primitives.

hyperbolix.nn_layers.HTCLinear

HTCLinear(
    in_features: int,
    out_features: int,
    *,
    rngs: Rngs,
    use_bias: bool = True,
    init_bound: float = 0.02,
    eps: float = 1e-07,
    param_dtype: DTypeLike = jnp.float32,
)

Bases: Module

Hyperbolic Transformation Component with learnable linear transformation.

This module wraps a Euclidean linear layer with the HTC operation, enabling learnable transformations between hyperboloid manifolds with different curvatures.

Parameters:

Name	Type	Description	Default
in_features	int	Input feature dimension (full hyperboloid dimension, including time component).	required
out_features	int	Output spatial dimension (time component is reconstructed automatically).	required
rngs	Rngs	Random number generators for parameter initialization.	required
use_bias	bool	Whether to include a bias term (default: True).	True
init_bound	float	Bound for uniform weight initialization. Weights are initialized from Uniform(-init_bound, init_bound). Small values keep initial outputs close to the hyperboloid origin for stable training (default: 0.02).	0.02
eps	float	Small value for numerical stability (default: 1e-7).	1e-07
param_dtype	DTypeLike	Storage dtype of the trainable parameters (default: jnp.float32). Compute precision of manifold operations is set by the manifold's dtype.	float32

Attributes:

Name	Type	Description
kernel	Param	Weight matrix of shape (in_features, out_features).
bias	Param or None	Bias vector of shape (out_features,) if use_bias=True, else None.
eps	float	Numerical stability parameter.

Notes

Weight Initialization: This layer uses small uniform initialization U(-0.02, 0.02) by default, matching the initialization used by FHNN/FHCNN layers. Standard deep learning initializations (Xavier, Lecun) produce weights that are too large for hyperbolic operations, causing gradient explosion and training instability.

See Also

hyperbolix.nn_layers.hyperboloid_core.htc : Core HTC function for functional transformations. HypLinearHyperboloidFHCNN : Alternative hyperbolic linear layer with sigmoid scaling.

References

Hypformer paper (citation to be added)

Examples:

>>> from flax import nnx
>>> from hyperbolix.nn_layers import HTCLinear
>>> from hyperbolix.manifolds import Hyperboloid
>>>
>>> # Create layer
>>> layer = HTCLinear(in_features=5, out_features=8, rngs=nnx.Rngs(0))
>>>
>>> # Forward pass
>>> manifold = Hyperboloid()
>>> x = jnp.ones((32, 5))  # batch of 32 points
>>> x = jax.vmap(manifold.proj, in_axes=(0, None))(x, 1.0)
>>> y = layer(x, c_in=1.0, c_out=2.0)
>>> y.shape
(32, 9)  # 8 spatial + 1 time

Source code in hyperbolix/nn_layers/hyperboloid_linear.py

def __init__(
    self,
    in_features: int,
    out_features: int,
    *,
    rngs: nnx.Rngs,
    use_bias: bool = True,
    init_bound: float = 0.02,
    eps: float = 1e-7,
    param_dtype: DTypeLike = jnp.float32,
):
    # Small uniform initialization for hyperbolic stability
    # Standard initializations (Lecun, Xavier) are too large and cause gradient explosion
    self.kernel = nnx.Param(
        jax.random.uniform(
            rngs.params(), (in_features, out_features), dtype=param_dtype, minval=-init_bound, maxval=init_bound
        )
    )
    if use_bias:
        self.bias = nnx.Param(jnp.zeros((out_features,), dtype=param_dtype))
    else:
        self.bias = None
    self.eps = eps

__call__

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

Apply HTC linear transformation.

Parameters:

Name	Type	Description	Default
x	Array of shape (batch, in_features)	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, out_features+1)	Output points on hyperboloid with curvature c_out.

Source code in hyperbolix/nn_layers/hyperboloid_linear.py

def __call__(
    self,
    x: Float[Array, "batch in_features"],
    c_in: float = 1.0,
    c_out: float = 1.0,
) -> Float[Array, "batch out_features_plus_1"]:
    """Apply HTC linear transformation.

    Parameters
    ----------
    x : Array of shape (batch, in_features)
        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, out_features+1)
        Output points on hyperboloid with curvature c_out.
    """

    def linear_fn(z):
        # Cast params to the working dtype so float64 weights (from global
        # jax_enable_x64) don't promote a float32 computation to float64.
        out = z @ self.kernel[...].astype(z.dtype)
        if self.bias is not None:
            out = out + self.bias[...].astype(z.dtype)
        return out

    return htc(x, linear_fn, c_in, c_out, self.eps)

hyperbolix.nn_layers.HypLinearHyperboloidFHCNN

HypLinearHyperboloidFHCNN(
    manifold_module: Manifold,
    in_dim: int,
    out_dim: int,
    *,
    rngs: Rngs,
    input_space: str = "manifold",
    init_scale: float = 2.3,
    learnable_scale: bool = False,
    eps: float = 1e-05,
    activation: Callable[[Array], Array] | None = None,
    normalize: bool = False,
    param_dtype: DTypeLike = jnp.float32,
)

Bases: Module

Fully Hyperbolic Convolutional Neural Networks fully connected layer (Hyperboloid model).

Computation steps: 0) Project the input tensor to the manifold (optional) 1) Apply activation (optional) 2) a) If normalize is True, compute the time and space coordinates of the output by applying a scaled sigmoid of the weight and biases transformed coordinates of the input or the result of the previous step. b) If normalize is False, compute the weight and biases transformed space coordinates of the input or the result of the previous step and set the time coordinate such that the result lies on the manifold.

Parameters:

Name	Type	Description	Default
manifold_module	object	Class-based Hyperboloid manifold instance	required
in_dim	int	Dimension of the input space	required
out_dim	int	Dimension of the output space	required
rngs	Rngs	Random number generators for parameter initialization	required
input_space	str	Type of the input tensor, either 'tangent' or 'manifold' (default: 'manifold'). Note: This is a static configuration - changing it after initialization requires recompilation.	'manifold'
init_scale	float	Initial value for the sigmoid scale parameter (default: 2.3)	2.3
learnable_scale	bool	Whether the scale parameter should be learnable (default: False)	False
eps	float	Small value to ensure that the time coordinate is bigger than 1/sqrt(c) (default: 1e-5)	1e-05
activation	callable or None	Activation function to apply before the linear transformation (default: None). Note: This is a static configuration - changing it after initialization requires recompilation.	None
normalize	bool	Whether to normalize the space coordinates before rescaling (default: False). Note: This is a static configuration - changing it after initialization requires recompilation.	False
param_dtype	DTypeLike	Storage dtype of the trainable parameters (default: jnp.float32). Compute precision of manifold operations is set by manifold.dtype.	float32

Notes

JIT Compatibility: This layer is designed to work with nnx.jit. Configuration parameters (input_space, activation, normalize) are treated as static and will be baked into the compiled function.

Relationship to HTC/HRC: When normalize=False and c_in = c_out, this layer uses the same time reconstruction pattern as htc: time = sqrt(||space||^2 + 1/c). The key difference is that FHCNN applies a linear transform to the full input and discards the computed time, while htc uses the linear output directly as spatial components. When normalize=True, FHCNN uses a learned sigmoid scaling which differs from both htc and hrc.

See Also

htc : Hyperbolic Transformation Component with curvature change support. Similar time reconstruction pattern when normalize=False. HTCLinear : Module wrapper for htc with learnable linear transformation.

References

Ahmad Bdeir, Kristian Schwethelm, and Niels Landwehr. "Fully hyperbolic convolutional neural networks for computer vision." arXiv preprint arXiv:2303.15919 (2023).

Source code in hyperbolix/nn_layers/hyperboloid_linear.py

def __init__(
    self,
    manifold_module: Manifold,
    in_dim: int,
    out_dim: int,
    *,
    rngs: nnx.Rngs,
    input_space: str = "manifold",
    init_scale: float = 2.3,
    learnable_scale: bool = False,
    eps: float = 1e-5,
    activation: Callable[[Array], Array] | None = None,
    normalize: bool = False,
    param_dtype: DTypeLike = jnp.float32,
):
    if input_space not in ["tangent", "manifold"]:
        raise ValueError(f"input_space must be either 'tangent' or 'manifold', got '{input_space}'")

    # Static configuration (treated as compile-time constants for JIT)
    validate_hyperboloid_manifold(manifold_module, required_methods=("expmap_0",))
    self.manifold = manifold_module
    self.in_dim = in_dim
    self.out_dim = out_dim
    self.input_space = input_space
    self.eps = eps
    self.activation = activation
    self.normalize = normalize

    # Trainable parameters
    bound = 0.02
    weight_init = jax.random.uniform(rngs.params(), (out_dim, in_dim), dtype=param_dtype, minval=-bound, maxval=bound)
    self.kernel = nnx.Param(weight_init)
    self.bias = nnx.Param(jnp.zeros((1, out_dim), dtype=param_dtype))

    # Scale parameter for sigmoid
    if learnable_scale:
        self.scale = nnx.Param(jnp.array(init_scale, dtype=param_dtype))
    else:
        # For non-learnable scale, store as regular Python float (static)
        self.scale = init_scale

__call__

__call__(
    x: Float[Array, "batch in_dim"], c: float = 1.0
) -> Float[Array, "batch out_dim"]

Forward pass through the FHCNN hyperbolic linear layer.

Parameters:

Name	Type	Description	Default
x	Array of shape (batch, in_dim)	Input tensor where the hyperbolic_axis is last. x.shape[-1] must equal self.in_dim.	required
c	float	Manifold curvature (default: 1.0)	1.0

Returns:

Name	Type	Description
res	Array of shape (batch, out_dim)	Output on the Hyperboloid manifold

Source code in hyperbolix/nn_layers/hyperboloid_linear.py

def __call__(
    self,
    x: Float[Array, "batch in_dim"],
    c: float = 1.0,
) -> Float[Array, "batch out_dim"]:
    """
    Forward pass through the FHCNN hyperbolic linear layer.

    Parameters
    ----------
    x : Array of shape (batch, in_dim)
        Input tensor where the hyperbolic_axis is last. x.shape[-1] must equal self.in_dim.
    c : float
        Manifold curvature (default: 1.0)

    Returns
    -------
    res : Array of shape (batch, out_dim)
        Output on the Hyperboloid manifold
    """
    scale_val = self.scale[...] if isinstance(self.scale, nnx.Param) else self.scale
    return _fhcnn_forward(
        x,
        self.kernel[...],
        self.bias[...],
        self.manifold,
        c,
        self.input_space,
        self.activation,
        self.normalize,
        scale_val,
        self.eps,
    )

hyperbolix.nn_layers.HypLinearHyperboloidFHNN

HypLinearHyperboloidFHNN(
    manifold_module: Manifold,
    in_dim: int,
    out_dim: int,
    *,
    rngs: Rngs,
    input_space: str = "manifold",
    init_scale: float = 2.3,
    eps: float = 1e-05,
    activation: Callable[[Array], Array] | None = None,
    dropout_rate: float | None = None,
    param_dtype: DTypeLike = jnp.float32,
)

Bases: Module

Fully Hyperbolic Neural Networks linear layer (Chen et al. 2021).

Time-primary parameterization: the time coordinate is computed via scaled sigmoid with an additive floor at 1/sqrt(c), and spatial coordinates are rescaled so the output lies on the hyperboloid.

Computation steps: 0) Project the input tensor to the manifold (optional) 1) Apply activation (optional) 2) Apply dropout (optional) 3) Linear transform: z = W @ x + b 4) Time: y0 = exp(scale) * sigmoid(z0) + 1/sqrt(c) + eps 5) Spatial: y_rem = sqrt(y0^2 - 1/c) / ||z_rem|| * z_rem 6) Output: [y0, y_rem] on the hyperboloid

Parameters:

Name	Type	Description	Default
manifold_module	object	Class-based Hyperboloid manifold instance	required
in_dim	int	Ambient input dimension (d+1, including time)	required
out_dim	int	Ambient output dimension (d+1, including time)	required
rngs	Rngs	Random number generators for parameter initialization	required
input_space	str	Type of the input tensor, either 'tangent' or 'manifold' (default: 'manifold'). Note: This is a static configuration - changing it after initialization requires recompilation.	'manifold'
init_scale	float	Initial value for the learnable sigmoid scale (default: 2.3)	2.3
eps	float	Numerical stability epsilon (default: 1e-5)	1e-05
activation	callable or None	Activation function to apply before the linear transformation (default: None). Note: This is a static configuration - changing it after initialization requires recompilation.	None
dropout_rate	float or None	Dropout rate applied before the linear transformation (default: None).	None
param_dtype	DTypeLike	Storage dtype of the trainable parameters (default: jnp.float32). Compute precision of manifold operations is set by manifold.dtype.	float32

Notes

JIT Compatibility: This layer is designed to work with nnx.jit. Configuration parameters (input_space, activation) are treated as static and will be baked into the compiled function.

Weight Initialization: Weights are initialized as tangent vectors at the hyperboloid origin: U(-0.02, 0.02) with the time column (column 0) zeroed. This matches the Chen et al. 2021 init.

Relationship to FHCNN: Both FHNN and FHCNN ensure outputs lie on the hyperboloid, but differ in which coordinate is primary. FHNN controls the time coordinate via sigmoid with an additive floor at 1/sqrt(c), then derives the spatial norm from the hyperboloid constraint. FHCNN (normalize=True) controls the spatial norm via sigmoid, then derives time.

See Also

HypLinearHyperboloidFHCNN : FHCNN layer with spatial-primary parameterization. HypLinearHyperboloidPLFC : PLFC layer using MLR + sinh diffeomorphism.

References

Weize Chen, et al. "Fully hyperbolic neural networks." arXiv preprint arXiv:2105.14686 (2021).

Source code in hyperbolix/nn_layers/hyperboloid_linear.py

def __init__(
    self,
    manifold_module: Manifold,
    in_dim: int,
    out_dim: int,
    *,
    rngs: nnx.Rngs,
    input_space: str = "manifold",
    init_scale: float = 2.3,
    eps: float = 1e-5,
    activation: Callable[[Array], Array] | None = None,
    dropout_rate: float | None = None,
    param_dtype: DTypeLike = jnp.float32,
):
    if input_space not in ["tangent", "manifold"]:
        raise ValueError(f"input_space must be either 'tangent' or 'manifold', got '{input_space}'")

    # Static configuration (treated as compile-time constants for JIT)
    validate_hyperboloid_manifold(manifold_module, required_methods=("expmap_0",))
    self.manifold = manifold_module
    self.in_dim = in_dim
    self.out_dim = out_dim
    self.input_space = input_space
    self.eps = eps
    self.activation = activation

    # FHNN weight init: U(-0.02, 0.02) with time column zeroed (tangent vectors at origin)
    bound = 0.02
    weight_init = jax.random.uniform(rngs.params(), (out_dim, in_dim), dtype=param_dtype, minval=-bound, maxval=bound)
    weight_init = weight_init.at[:, 0].set(0.0)
    self.kernel = nnx.Param(weight_init)
    self.bias = nnx.Param(jnp.zeros((1, out_dim), dtype=param_dtype))

    # Learnable scale for the sigmoid (always learnable in FHNN)
    self.scale = nnx.Param(jnp.array(init_scale, dtype=param_dtype))

    # Optional dropout
    if dropout_rate is not None and dropout_rate > 0:
        self.dropout = nnx.Dropout(dropout_rate, rngs=rngs)
    else:
        self.dropout = None

__call__

__call__(
    x: Float[Array, "batch in_dim"],
    c: float = 1.0,
    deterministic: bool = True,
) -> Float[Array, "batch out_dim"]

Forward pass through the FHNN hyperbolic linear layer.

Parameters:

Name	Type	Description	Default
x	Array of shape (batch, in_dim)	Input tensor. x.shape[-1] must equal self.in_dim.	required
c	float	Manifold curvature (default: 1.0)	1.0
deterministic	bool	If True, dropout is disabled (default: True).	True

Returns:

Name	Type	Description
res	Array of shape (batch, out_dim)	Output on the Hyperboloid manifold

Source code in hyperbolix/nn_layers/hyperboloid_linear.py

def __call__(
    self,
    x: Float[Array, "batch in_dim"],
    c: float = 1.0,
    deterministic: bool = True,
) -> Float[Array, "batch out_dim"]:
    """Forward pass through the FHNN hyperbolic linear layer.

    Parameters
    ----------
    x : Array of shape (batch, in_dim)
        Input tensor. x.shape[-1] must equal self.in_dim.
    c : float
        Manifold curvature (default: 1.0)
    deterministic : bool
        If True, dropout is disabled (default: True).

    Returns
    -------
    res : Array of shape (batch, out_dim)
        Output on the Hyperboloid manifold
    """
    # Build dropout closure for the pure function
    dropout_module = self.dropout
    if dropout_module is not None:
        dropout_fn = lambda z: dropout_module(z, deterministic=deterministic)  # noqa: E731
    else:
        dropout_fn = None

    return _fhnn_forward(
        x,
        self.kernel[...],
        self.bias[...],
        self.manifold,
        c,
        self.input_space,
        self.activation,
        dropout_fn,
        self.scale[...],
        self.eps,
    )

hyperbolix.nn_layers.HypLinearHyperboloidPLFC

HypLinearHyperboloidPLFC(
    manifold_module: Hyperboloid,
    in_dim: int,
    out_dim: int,
    *,
    rngs: Rngs,
    input_space: str = "manifold",
    clamping_factor: float = 1.0,
    smoothing_factor: float = 50.0,
    v_max: float = 10.0,
    use_gyro_bias: bool = False,
    kernel_init_std: float = 0.02,
    param_dtype: DTypeLike = jnp.float32,
)

Bases: Module

Point-to-hyperplane Lorentz fully connected (PLFC) layer (Hyperboloid model).

Each output coordinate is the signed Lorentz distance from the input point to a learned Lorentz hyperplane (the Lorentz MLR score). The output point is recovered in closed form via the element-wise sinh diffeomorphism and the hyperboloid time constraint, so the signed distance from the output to the k-th coordinate hyperplane equals the k-th score (margin preservation, Shi et al. 2026, Thm. 1-2). This is the Lorentz analog of the HNN++ point-to-hyperplane FC formulation (Shimizu et al. 2020). Formerly named HypLinearHyperboloidPP.

Computation steps: 0) Project the input tensor onto the manifold (optional) 1) Compute the multinomial linear regression score(s) via compute_mlr 2) Hard-clamp the scaled scores sqrt(c)*v to ±v_max (output-side guard) 3) Apply element-wise sinh diffeomorphism to obtain spatial coordinates 4) Reconstruct time coordinate from the hyperboloid constraint 5) Optionally add an intrinsic gyro-bias y ← y ⊕ exp_0([0, b])

Parameters:

Name	Type	Description	Default
manifold_module	object	Class-based Hyperboloid manifold instance	required
in_dim	int	Full input dimension (ambient, d+1)	required
out_dim	int	Full output dimension (ambient, d+1)	required
rngs	Rngs	Random number generators for parameter initialization	required
input_space	str	Type of the input tensor, either 'tangent' or 'manifold' (default: 'manifold'). Note: This is a static configuration - changing it after initialization requires recompilation.	'manifold'
clamping_factor	float	Clamping factor for the multinomial linear regression output (default: 1.0)	1.0
smoothing_factor	float	Smoothing factor for the multinomial linear regression output (default: 50.0)	50.0
v_max	float	Output-side guard: the sinh argument sqrt(c)*v is hard-clipped to ±v_max, bounding the output spatial norm by sinh(v_max)/sqrt(c) (default: 10.0, matching the Shi et al. 2026 reference implementation).	10.0
use_gyro_bias	bool	If True, add a learnable intrinsic bias via Lorentz gyroaddition, y ← y ⊕ exp_0([0, b]) with b initialized to zero — the gyrogroup identity, so the bias is a no-op at initialization (default: False).	False
kernel_init_std	float	Standard deviation of the normal kernel init (default: 0.02, the Shi et al. 2026 PLFC reference value). Use 1.0 to recover the previous HNN++-style init (Shimizu et al. 2020); note that large kernels push the pre-guard scores toward the v_max saturation regime.	0.02
param_dtype	DTypeLike	Storage dtype of the trainable parameters (default: jnp.float32). Compute precision of manifold operations is set by manifold.dtype.	float32

Notes

JIT Compatibility: This layer is designed to work with nnx.jit. Configuration parameters (input_space, clamping_factor, smoothing_factor, v_max) are treated as static and will be baked into the compiled function.

References

Xianglong Shi, Ziheng Chen, Yunhan Jiang, and Nicu Sebe. "Intrinsic Lorentz Neural Network." ICLR 2026, Sec. 4.1 (PLFC layer, Theorem 1). Shimizu Ryohei, Yusuke Mukuta, and Tatsuya Harada. "Hyperbolic neural networks++." arXiv preprint arXiv:2006.08210 (2020).

Source code in hyperbolix/nn_layers/hyperboloid_linear.py

def __init__(
    self,
    manifold_module: Hyperboloid,
    in_dim: int,
    out_dim: int,
    *,
    rngs: nnx.Rngs,
    input_space: str = "manifold",
    clamping_factor: float = 1.0,
    smoothing_factor: float = 50.0,
    v_max: float = 10.0,
    use_gyro_bias: bool = False,
    kernel_init_std: float = 0.02,
    param_dtype: DTypeLike = jnp.float32,
):
    if input_space not in ["tangent", "manifold"]:
        raise ValueError(f"input_space must be either 'tangent' or 'manifold', got '{input_space}'")

    # Static configuration (treated as compile-time constants for JIT)
    required_methods = ("expmap_0", "compute_mlr")
    if use_gyro_bias:
        required_methods = ("expmap_0", "compute_mlr", "embed_spatial_0", "addition")
    validate_hyperboloid_manifold(manifold_module, required_methods=required_methods)
    self.manifold = manifold_module
    self.in_dim = in_dim
    self.out_dim = out_dim
    self.input_space = input_space
    self.clamping_factor = clamping_factor
    self.smoothing_factor = smoothing_factor
    _assert_v_max_safe(v_max)
    self.v_max = v_max

    # Trainable parameters — small normal init (Shi et al. 2026 PLFC reference)
    in_spatial = in_dim - 1
    out_spatial = out_dim - 1
    kernel_init = kernel_init_std * jax.random.normal(rngs.params(), (out_spatial, in_spatial), dtype=param_dtype)
    self.kernel = nnx.Param(kernel_init)
    self.bias = nnx.Param(jnp.zeros((out_spatial, 1), dtype=param_dtype))

    # Gyro-bias: spatial tangent vector at the origin, zero-initialized
    if use_gyro_bias:
        self.gyro_bias = nnx.Param(jnp.zeros((out_spatial,), dtype=param_dtype))
    else:
        self.gyro_bias = None

__call__

__call__(
    x: Float[Array, "batch in_dim"], c: float = 1.0
) -> Float[Array, "batch out_dim"]

Forward pass through the PLFC hyperboloid linear layer.

Parameters:

Name	Type	Description	Default
x	Array of shape (batch, in_dim)	Input tensor. x.shape[-1] must equal self.in_dim.	required
c	float	Manifold curvature (default: 1.0)	1.0

Returns:

Name	Type	Description
res	Array of shape (batch, out_dim)	Output on the Hyperboloid manifold

Source code in hyperbolix/nn_layers/hyperboloid_linear.py

def __call__(
    self,
    x: Float[Array, "batch in_dim"],
    c: float = 1.0,
) -> Float[Array, "batch out_dim"]:
    """
    Forward pass through the PLFC hyperboloid linear layer.

    Parameters
    ----------
    x : Array of shape (batch, in_dim)
        Input tensor. x.shape[-1] must equal self.in_dim.
    c : float
        Manifold curvature (default: 1.0)

    Returns
    -------
    res : Array of shape (batch, out_dim)
        Output on the Hyperboloid manifold
    """
    return _hyperboloid_plfc_forward(
        x,
        self.kernel[...],
        self.bias[...],
        self.gyro_bias[...] if self.gyro_bias is not None else None,
        self.manifold,
        c,
        self.input_space,
        self.clamping_factor,
        self.smoothing_factor,
        self.v_max,
    )

hyperbolix.nn_layers.HypLinearHyperboloidBusemann

HypLinearHyperboloidBusemann(
    manifold_module: Hyperboloid,
    in_dim: int,
    out_dim: int,
    *,
    rngs: Rngs,
    input_space: str = "manifold",
    activation: Callable[[Array], Array] | None = None,
    v_max: float = 10.0,
    use_gyro_bias: bool = False,
    param_dtype: DTypeLike = jnp.float32,
)

Bases: Module

Busemann fully connected (BFC) layer (Hyperboloid / Lorentz model).

Maps a hyperboloid point to a hyperboloid point (Chen et al. 2026, Thm. 4.2). Each output spatial coordinate is the activated Busemann logit φ(-alpha_k·B^{v_k}(x) + b_k); the point is recovered via the shared sinh_lift_to_hyperboloid output map (the same closed form as HypLinearHyperboloidPLFC, but with a point-to-horosphere score instead of point-to-hyperplane).

Parameters:

Name	Type	Description	Default
manifold_module	Hyperboloid	Class-based Hyperboloid manifold instance.	required
in_dim	int	Input ambient dimension (d + 1, time included).	required
out_dim	int	Output ambient dimension (d_out + 1, time included).	required
rngs	Rngs	Random number generators for parameter initialization.	required
input_space	str	"manifold" (default) or "tangent" (lift via expmap_0 first). Static for JIT.	'manifold'
activation	Callable or None	Optional Euclidean activation φ applied to the Busemann logits (default: identity). Avoid relu when stacking several of these layers on high-dimensional input: at a random weight-normalized direction init, the Busemann score is positive for all but a narrow cone of directions, so relu can zero every output unit at once and collapse the whole layer to the manifold origin — a fixed point that self-reinforces through later layers. tanh avoids this by never fully zeroing the logit.	None
v_max	float	Output-side guard: √c·φ(u) is hard-clipped to ±v_max before the sinh diffeomorphism (default: 10.0, the Shi et al. 2026 reference value).	10.0
use_gyro_bias	bool	If True, add a learnable intrinsic bias y ← y ⊕ exp_0([0, b]) via Lorentz gyroaddition, b zero-initialized (gyrogroup identity ⇒ no-op at init) (default: False).	False
param_dtype	DTypeLike	Storage dtype of the trainable parameters (default: jnp.float32).	float32

References

Chen, Schölkopf, and Sebe. "Hyperbolic Busemann Neural Networks." 2026, Thm. 4.2.

Source code in hyperbolix/nn_layers/busemann_linear.py

def __init__(
    self,
    manifold_module: Hyperboloid,
    in_dim: int,
    out_dim: int,
    *,
    rngs: nnx.Rngs,
    input_space: str = "manifold",
    activation: Callable[[Array], Array] | None = None,
    v_max: float = 10.0,
    use_gyro_bias: bool = False,
    param_dtype: DTypeLike = jnp.float32,
):
    if input_space not in ["tangent", "manifold"]:
        raise ValueError(f"input_space must be either 'tangent' or 'manifold', got '{input_space}'")

    required_methods = ("expmap_0", "busemann")
    if use_gyro_bias:
        required_methods = ("expmap_0", "busemann", "embed_spatial_0", "addition")
    validate_hyperboloid_manifold(manifold_module, required_methods=required_methods)
    _assert_v_max_safe(v_max)

    self.manifold = manifold_module
    self.in_dim = in_dim
    self.out_dim = out_dim
    self.input_space = input_space
    self.activation = activation
    self.v_max = v_max

    in_spatial = in_dim - 1
    out_spatial = out_dim - 1
    # van Spengler / reference BFC init, computed from the spatial dims.
    self.kernel, self.log_scale, self.bias = init_weight_norm_params(
        rngs, out_spatial, in_spatial, std=(2 * in_spatial * out_spatial) ** -0.5, param_dtype=param_dtype
    )

    if use_gyro_bias:
        self.gyro_bias = nnx.Param(jnp.zeros((out_spatial,), dtype=param_dtype))
    else:
        self.gyro_bias = None

__call__

__call__(
    x: Float[Array, "batch in_dim"], c: float = 1.0
) -> Float[Array, "batch out_dim"]

Forward pass returning a hyperboloid point, shape (batch, out_dim).

Source code in hyperbolix/nn_layers/busemann_linear.py

def __call__(
    self,
    x: Float[Array, "batch in_dim"],
    c: float = 1.0,
) -> Float[Array, "batch out_dim"]:
    """Forward pass returning a hyperboloid point, shape (batch, out_dim)."""
    u_BO = busemann_score(self.manifold, x, self.kernel[...], self.log_scale[...], self.bias[...], c, self.input_space)
    if self.activation is not None:
        u_BO = self.activation(u_BO)
    y_BAo = sinh_lift_to_hyperboloid(u_BO, c, self.v_max)

    if self.gyro_bias is not None:
        # Paper notation F(x) ⊕ b (Sec. 4.2), matching HypLinearHyperboloidPLFC's gyro-bias.
        # (The reference code applies b ⊕ F(x); the two coincide at the zero-init identity.)
        bias_tangent_Ao = self.manifold.embed_spatial_0(self.gyro_bias[...].astype(y_BAo.dtype))
        bias_point_Ao = self.manifold.expmap_0(bias_tangent_Ao, c)
        y_BAo = jax.vmap(self.manifold.addition, in_axes=(0, None, None))(y_BAo, bias_point_Ao, c)

    return y_BAo

FGG Linear (Klis et al. 2026)

FGGLinear achieves linear (not logarithmic) growth of hyperbolic distance via the sinh/arcsinh cancellation in the Lorentzian activation chain. It defaults to a norm-preserving fan_out init suited to unnormalized stacks — see the init-scales note.

hyperbolix.nn_layers.FGGLinear

FGGLinear(
    in_features: int,
    out_features: int,
    *,
    rngs: Rngs,
    activation: Callable[[Array], Array] | None = None,
    reset_params: str = "fan_out",
    use_weight_norm: bool = False,
    init_bias: float = 0.0,
    gain: float = 1.0,
    eps: float = 1e-07,
    param_dtype: DTypeLike = jnp.float32,
)

Bases: Module

Fast and Geometrically Grounded Lorentz linear layer.

Implements the FGG linear layer from Klis et al. 2026. The key insight is that the sinh/arcsinh cancellation in the Lorentzian activation chain simplifies the forward pass to: matmul with spacelike V matrix -> Euclidean activation -> time reconstruction. This achieves linear growth of hyperbolic distance (vs logarithmic for Chen et al. 2022) and ~3x faster training/inference.

Forward pass: 1. Build spacelike V matrix from (U, b) with Minkowski metric absorbed 2. z = x @ V (Minkowski inner products via a single matmul) 3. z = h(z) (Euclidean activation, e.g. ReLU) 4. y_0 = sqrt(||z||^2 + 1/c) (time reconstruction) 5. y = [y_0, z] (on hyperboloid)

Parameters:

Name	Type	Description	Default
in_features	int	Input ambient dimension (D_in + 1), including time component.	required
out_features	int	Output ambient dimension (D_out + 1), including time component.	required
rngs	Rngs	Random number generators for parameter initialization.	required
activation	Callable or None	Euclidean activation function applied after matmul (default: None).	None
reset_params	str	Weight initialization scheme: "eye", "xavier", "kaiming", "lorentz_kaiming", "fan_out", or "mlr" (default: "fan_out"). The "fan_out" default (std sqrt(1/out_spatial)) is norm-preserving (||z|| ~= gain * ||x_spatial||), a deliberate deviation from the Klis et al. 2026 classification reference: it suits unnormalized stacks (e.g. an RL backbone feeding a bounded Poincare-ball projection), where the reference's BatchNorm regime does not apply. Pass reset_params="eye", init_bias=0.5 to recover the previous reference-style init.	'fan_out'
use_weight_norm	bool	If True, reparameterize U as g * v / ||v|| for weight normalization (default: False). Note that gain and the "fan_out" scale are renormalized away in this mode (the magnitude is set by kernel_scale).	False
init_bias	float	Initial value for bias entries (default: 0.0). A zero bias removes the ~sqrt(out) * init_bias quadrature term that build_spacelike_V injects into the time row (-||w|| * sinh(-sqrt(c) * b / ||w||)); pair with "fan_out" for norm preservation.	0.0
gain	float	Multiplier on the random init (default: 1.0). With reset_params="fan_out" it sets the effective column std to gain / sqrt(out_spatial). No-op for "eye"; renormalized away under use_weight_norm=True.	1.0
eps	float	Numerical stability floor (default: 1e-7).	1e-07
param_dtype	DTypeLike	Storage dtype of the trainable parameters (default: jnp.float32). Compute precision of manifold operations is set by the manifold's dtype.	float32

References

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

Examples:

>>> from flax import nnx
>>> from hyperbolix.nn_layers import FGGLinear
>>> import jax.numpy as jnp
>>>
>>> layer = FGGLinear(33, 65, rngs=nnx.Rngs(0), activation=jax.nn.relu)
>>> x = jnp.ones((8, 33))
>>> # project to hyperboloid
>>> x = x.at[:, 0].set(jnp.sqrt(jnp.sum(x[:, 1:]**2, axis=-1) + 1.0))
>>> y = layer(x, c=1.0)
>>> y.shape
(8, 65)

Source code in hyperbolix/nn_layers/hyperboloid_linear.py

def __init__(
    self,
    in_features: int,
    out_features: int,
    *,
    rngs: nnx.Rngs,
    activation: Callable[[jax.Array], jax.Array] | None = None,
    reset_params: str = "fan_out",
    use_weight_norm: bool = False,
    init_bias: float = 0.0,
    gain: float = 1.0,
    eps: float = 1e-7,
    param_dtype: DTypeLike = jnp.float32,
):
    in_spatial = in_features - 1  # I
    out_spatial = out_features - 1  # O

    self.in_features = in_features
    self.out_features = out_features
    self.activation = activation
    self.use_weight_norm = use_weight_norm
    self.eps = eps

    # Euclidean weight U (I, O). Default "fan_out" std sqrt(1/out_spatial) is
    # norm-preserving; gain scales the random init (no-op for "eye").
    U_init = _fgg_weight_init(
        rngs.params(), in_spatial, out_spatial, in_features, out_features, reset_params, param_dtype, gain=gain
    )

    # Weight normalization: decompose kernel = softplus(kernel_scale) * kernel_dir / ||kernel_dir||
    if use_weight_norm:
        # Reference: kernel_dir from reset_params (normalized in forward), kernel_scale fixed magnitude
        self.kernel_dir = nnx.Param(U_init)  # (I, O) direction
        g_init_val = jnp.sqrt(1.0 / (in_features + out_features))
        self.kernel_scale = nnx.Param(jnp.full((out_spatial,), g_init_val, dtype=param_dtype))  # (O,)
    else:
        self.kernel = nnx.Param(U_init)  # (I, O)

    # Bias: init to init_bias
    self.bias = nnx.Param(jnp.full((out_spatial,), init_bias, dtype=param_dtype))  # (O,)

__call__

__call__(
    x_BAi: Float[Array, "batch in_features"], c: float = 1.0
) -> Float[Array, "batch out_features"]

Forward pass through the FGG linear layer.

Parameters:

Name	Type	Description	Default
x_BAi	(Array, shape(B, Ai))	Input points on the hyperboloid with curvature c. Ai = in_features (ambient dimension).	required
c	float	Curvature parameter (default: 1.0).	1.0

Returns:

Name	Type	Description
y_BAo	(Array, shape(B, Ao))	Output points on the hyperboloid with curvature c. Ao = out_features (ambient dimension).

Source code in hyperbolix/nn_layers/hyperboloid_linear.py

def __call__(
    self,
    x_BAi: Float[Array, "batch in_features"],
    c: float = 1.0,
) -> Float[Array, "batch out_features"]:
    """Forward pass through the FGG linear layer.

    Parameters
    ----------
    x_BAi : Array, shape (B, Ai)
        Input points on the hyperboloid with curvature ``c``.
        Ai = in_features (ambient dimension).
    c : float, optional
        Curvature parameter (default: 1.0).

    Returns
    -------
    y_BAo : Array, shape (B, Ao)
        Output points on the hyperboloid with curvature ``c``.
        Ao = out_features (ambient dimension).
    """
    U_IO = self._get_kernel()  # (I, O)
    return _fgg_linear_forward(x_BAi, U_IO, self.bias[...], c, self.activation, self.eps)

Proper Velocity

HypLinearPV implements the PV fully-connected layer from Chen et al. 2026 (Thm 5.3 / Eq. 22): \(y_k = (1/\sqrt{c}) \cdot \sinh(\sqrt{c} \cdot v_k(x))\) where \(v_k(x)\) is the PV multinomial-logistic-regression signed margin. PV is an unconstrained \(\mathbb{R}^n\) model, so inputs and outputs are Euclidean with no projection step.

hyperbolix.nn_layers.HypLinearPV

HypLinearPV(
    manifold_module: ProperVelocity,
    in_dim: int,
    out_dim: int,
    *,
    rngs: Rngs,
    input_space: str = "manifold",
    inner_activation: Callable[[Array], Array]
    | None = None,
    clamping_factor: float = 1.0,
    smoothing_factor: float = 50.0,
    kernel_init_std: float | None = None,
    param_dtype: DTypeLike = jnp.float32,
)

Bases: Module

Hyperbolic Neural Networks fully connected layer (Proper Velocity model).

Implements Chen et al. 2026, Thm 5.3 / Eq. 22:

y_k = (1/√c) · sinh(√c · v_k(x))

where v_k(x) is the PV multinomial-logistic-regression signed margin from Eq. 19 (implemented as ProperVelocity.compute_mlr). The optional inner_activation realizes the Eq. 23 ablation from the paper.

Parameters:

Name	Type	Description	Default
manifold_module	ProperVelocity	Class-based Proper Velocity manifold instance.	required
in_dim	int	Dimension of the input space.	required
out_dim	int	Dimension of the output space.	required
rngs	Rngs	Random number generators for parameter initialization.	required
input_space	str	Type of the input tensor, either 'tangent' or 'manifold' (default: 'manifold'). Note: This is a static configuration — changing it after initialization requires recompilation.	'manifold'
inner_activation	Callable[[Array], Array] | None	Optional activation applied to v before the outer sinh (paper Eq. 23 ablation). Default None.	None
clamping_factor	float	Clamping factor for the MLR output (default: 1.0).	1.0
smoothing_factor	float	Smoothing factor for the MLR output (default: 50.0).	50.0
kernel_init_std	float | None	Standard deviation for the Gaussian kernel init. If None (default), uses He scaling sqrt(2 / in_dim) so stacked layers preserve variance under ReLU activations (for small MLR arguments the PV output reduces to an Euclidean linear map, so standard He analysis applies). Pass a specific value (e.g. 1e-2) when using this layer as the final classification/regression layer directly before a softmax/MSE loss — that matches the paper's PVManifoldMLR.reset_parameters recipe, which is tuned for receiving O(1)-variance features.	None
param_dtype	DTypeLike	Storage dtype of the trainable parameters (default: jnp.float32). Compute precision of manifold operations is set by manifold.dtype.	float32

Notes

JIT Compatibility: Configuration parameters (input_space, clamping_factor, smoothing_factor, inner_activation) are treated as static and are baked into the compiled function.

References

Chen et al. "Proper Velocity Neural Networks." ICLR 2026.

Source code in hyperbolix/nn_layers/pv_linear.py

def __init__(
    self,
    manifold_module: ProperVelocity,
    in_dim: int,
    out_dim: int,
    *,
    rngs: nnx.Rngs,
    input_space: str = "manifold",
    inner_activation: Callable[[Array], Array] | None = None,
    clamping_factor: float = 1.0,
    smoothing_factor: float = 50.0,
    kernel_init_std: float | None = None,
    param_dtype: DTypeLike = jnp.float32,
):
    if input_space not in ["tangent", "manifold"]:
        raise ValueError(f"input_space must be either 'tangent' or 'manifold', got '{input_space}'")

    validate_pv_manifold(
        manifold_module,
        required_methods=("expmap_0", "compute_mlr"),
    )
    self.manifold = manifold_module
    self.in_dim = in_dim
    self.out_dim = out_dim
    self.input_space = input_space
    self.inner_activation = inner_activation
    self.clamping_factor = clamping_factor
    self.smoothing_factor = smoothing_factor

    # Kernel init: He(in_dim) by default so deep stacks preserve variance
    # under ReLU. Override via ``kernel_init_std`` when using as a final
    # classification layer — see the ``kernel_init_std`` docstring above.
    # Bias init is uniform U(-1e-3, 1e-3), matching the paper reference.
    std = (2.0 / in_dim) ** 0.5 if kernel_init_std is None else kernel_init_std
    self.kernel = nnx.Param(jax.random.normal(rngs.params(), (out_dim, in_dim), dtype=param_dtype) * std)
    self.bias = nnx.Param(jax.random.uniform(rngs.params(), (out_dim, 1), dtype=param_dtype, minval=-1e-3, maxval=1e-3))

__call__

__call__(
    x: Float[Array, "batch in_dim"], c: float = 1.0
) -> Float[Array, "batch out_dim"]

Forward pass through the PV linear layer.

Parameters:

Name	Type	Description	Default
x	Array of shape (batch, in_dim)	Input — PV manifold points or tangent vectors at origin, depending on input_space.	required
c	float	Manifold curvature (default: 1.0).	1.0

Returns:

Name	Type	Description
res	Array of shape (batch, out_dim)	Output points on the PV manifold.

Source code in hyperbolix/nn_layers/pv_linear.py

def __call__(
    self,
    x: Float[Array, "batch in_dim"],
    c: float = 1.0,
) -> Float[Array, "batch out_dim"]:
    """
    Forward pass through the PV linear layer.

    Parameters
    ----------
    x : Array of shape (batch, in_dim)
        Input — PV manifold points or tangent vectors at origin, depending on
        ``input_space``.
    c : float
        Manifold curvature (default: 1.0).

    Returns
    -------
    res : Array of shape (batch, out_dim)
        Output points on the PV manifold.
    """
    return _pv_fc_forward(
        x,
        self.kernel[...],
        self.bias[...],
        self.manifold,
        c,
        self.input_space,
        self.clamping_factor,
        self.smoothing_factor,
        self.inner_activation,
    )

Example

import jax
from flax import nnx
from hyperbolix.nn_layers import HypLinearPoincare
from hyperbolix.manifolds import Poincare

poincare = Poincare()
layer = HypLinearPoincare(manifold_module=poincare, in_dim=32, out_dim=16, rngs=nnx.Rngs(0))

x = jax.random.normal(jax.random.PRNGKey(1), (10, 32)) * 0.3
x_proj = jax.vmap(poincare.proj, in_axes=(0, None))(x, 1.0)
output = layer(x_proj, c=1.0)
print(output.shape)  # (10, 16)

See the NN Layers guide for multi-layer composition patterns (HTC classifier, hybrid CNN + Poincaré head).