Regression & MLR

Single-layer hyperbolic classifiers (multinomial logistic regression heads). For which head to pair with which backbone, see the NN Layers guide.

Two decision geometries are available: point-to-hyperplane heads (Ganea/Shimizu/Bdeir lineage — the HypRegression* and FGGLorentzMLR layers) and point-to-horosphere Busemann heads (Chen et al. 2026), whose logits are \(u_k(x) = -\alpha_k B^{v_k}(x) + b_k\) from the closed-form Busemann function (Hyperboloid.busemann / Poincare.busemann).

Poincaré

hyperbolix.nn_layers.HypRegressionPoincare

HypRegressionPoincare(
    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,
    curvature: float | Callable[[], Any] = 1.0,
    param_dtype: DTypeLike = jnp.float32,
)

Bases: Module

Hyperbolic Neural Networks multinomial linear regression layer (Poincaré ball model).

Computation steps: 0) Project the input tensor onto the manifold (optional) 1) Compute the multinomial linear 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
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, clamping_factor, smoothing_factor) are treated as static and will be baked into the compiled function.

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_regression.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,
    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", "ptransp_0", "conformal_factor", "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 (Euclidean)
    self.kernel = nnx.Param(jax.random.normal(rngs.params(), (out_dim, in_dim), dtype=param_dtype))
    # Manifold bias (initialized to small random values)
    self.bias = ManifoldParam(
        jax.random.normal(rngs.params(), (out_dim, in_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 regression 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)	Multinomial linear regression scores

Source code in hyperbolix/nn_layers/poincare_regression.py

def __call__(
    self,
    x: Float[Array, "batch in_dim"],
    c: float = 1.0,
) -> Float[Array, "batch out_dim"]:
    """
    Forward pass through the hyperbolic regression 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)
        Multinomial linear regression scores
    """
    # Map to manifold if needed (static branch - JIT friendly)
    if self.input_space == "tangent":
        x = jax.vmap(self.manifold.expmap_0, in_axes=(0, None), out_axes=0)(x, c)

    # Project bias to manifold (vmap over out_dim dimension)
    bias_PD = jax.vmap(self.manifold.proj, in_axes=(0, None), out_axes=0)(self.bias[...], c)  # type: ignore[arg-type]

    # Parallel transport weight from tangent space at origin to tangent space at bias
    pt_weight_PD = jax.vmap(self.manifold.ptransp_0, in_axes=(0, 0, None), out_axes=0)(self.kernel[...], bias_PD, c)

    # Compute the multinomial linear regression score(s)
    res_BP = self._compute_mlr(x, pt_weight_PD, bias_PD, c)
    return res_BP

hyperbolix.nn_layers.HypRegressionPoincarePP

HypRegressionPoincarePP(
    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 ++ multinomial linear regression layer (Poincaré ball model).

Computation steps: 0) Project the input tensor onto the manifold (optional) 1) Compute the multinomial linear 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_regression.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", "ptransp_0", "conformal_factor", "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 — scaled to match reference (van Spengler et al. 2023)
    # Reference uses std = (2 * in_dim * out_dim)^{-0.5}; unscaled normal(0,1) gives
    # row norms ≈ sqrt(in_dim) which overwhelms the MLR output scaling.
    std = 1.0 / jnp.sqrt(2.0 * in_dim * out_dim)
    self.kernel = nnx.Param(jax.random.normal(rngs.params(), (out_dim, in_dim), dtype=param_dtype) * std)
    # Scalar bias (initialized to small random values)
    self.bias = nnx.Param(jax.random.normal(rngs.params(), (out_dim, 1), dtype=param_dtype) * 0.01)

__call__

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

Forward pass through the HNN++ hyperbolic regression 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)	Multinomial linear regression scores

Source code in hyperbolix/nn_layers/poincare_regression.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 regression 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)
        Multinomial linear regression scores
    """
    # Map to manifold if needed (static branch - JIT friendly)
    if self.input_space == "tangent":
        x = jax.vmap(self.manifold.expmap_0, in_axes=(0, None), out_axes=0)(x, c)

    # Compute multinomial linear regression
    res = self.manifold.compute_mlr_pp(
        x,
        self.kernel[...],
        self.bias[...],
        c,
        self.clamping_factor,
        self.smoothing_factor,
    )

    return res

hyperbolix.nn_layers.HypRegressionPoincareBusemann

HypRegressionPoincareBusemann(
    manifold_module: Poincare,
    in_dim: int,
    out_dim: int,
    *,
    rngs: Rngs,
    input_space: str = "manifold",
    param_dtype: DTypeLike = jnp.float32,
)

Bases: Module

Busemann MLR classification head (Poincaré ball model).

Computes per-class logits u_k(x) = -alpha_k·B^{v_k}(x) + b_k from the Poincaré Busemann function (point-to-horosphere distance), returning Euclidean logits — not manifold points.

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	Number of classes.	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'
param_dtype	DTypeLike	Storage dtype of the trainable parameters (default: jnp.float32). Compute precision of manifold operations is set by manifold.dtype.	float32

References

Chen, Schölkopf, and Sebe. "Hyperbolic Busemann Neural Networks." 2026, Sec. 3.

Source code in hyperbolix/nn_layers/busemann_regression.py

def __init__(
    self,
    manifold_module: Poincare,
    in_dim: int,
    out_dim: int,
    *,
    rngs: nnx.Rngs,
    input_space: str = "manifold",
    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_poincare_manifold(manifold_module, required_methods=("expmap_0", "busemann"))
    self.manifold = manifold_module
    self.in_dim = in_dim
    self.out_dim = out_dim
    self.input_space = input_space

    # Direction dim equals the spatial dim (no time component on the ball).
    in_spatial = in_dim
    self.kernel, self.log_scale, self.bias = init_weight_norm_params(
        rngs, out_dim, in_spatial, std=in_spatial**-0.5, param_dtype=param_dtype
    )

__call__

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

Forward pass returning Euclidean Busemann logits, shape (batch, out_dim).

Source code in hyperbolix/nn_layers/busemann_regression.py

def __call__(
    self,
    x: Float[Array, "batch in_dim"],
    c: float = 1.0,
) -> Float[Array, "batch out_dim"]:
    """Forward pass returning Euclidean Busemann logits, shape (batch, out_dim)."""
    return busemann_score(self.manifold, x, self.kernel[...], self.log_scale[...], self.bias[...], c, self.input_space)

Hyperboloid

FGGLorentzMLR is the FGG-family head (Klis et al. 2026); pair it with an FGG linear stack.

hyperbolix.nn_layers.HypRegressionHyperboloid

HypRegressionHyperboloid(
    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,
    param_dtype: DTypeLike = jnp.float32,
)

Bases: Module

Fully Hyperbolic Convolutional Neural Networks multinomial linear regression layer (Hyperboloid model).

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

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'
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

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_regression.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,
    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", "compute_mlr"))
    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
    # kernel lies in the tangent space of the Hyperboloid origin, so the time coordinate along axis is zero
    # and the true fan-in is the spatial dimension (in_dim - 1), not in_dim. Scaled to match reference
    # (van Spengler et al. 2023): std = (2 * in_spatial * out_dim)^{-0.5}; unscaled normal(0,1) gives row
    # norms ≈ sqrt(in_dim - 1) which overwhelms the MLR output scaling.
    in_spatial = in_dim - 1
    std = 1.0 / jnp.sqrt(2.0 * in_spatial * out_dim)
    self.kernel = nnx.Param(jax.random.normal(rngs.params(), (out_dim, in_spatial), dtype=param_dtype) * std)
    # Scalar bias (initialized to small random values)
    self.bias = nnx.Param(jax.random.normal(rngs.params(), (out_dim, 1), dtype=param_dtype) * 0.01)

__call__

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

Forward pass through the hyperbolic regression 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)	Multinomial linear regression scores

Source code in hyperbolix/nn_layers/hyperboloid_regression.py

def __call__(
    self,
    x: Float[Array, "batch in_dim"],
    c: float = 1.0,
) -> Float[Array, "batch out_dim"]:
    """
    Forward pass through the hyperbolic regression 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)
        Multinomial linear regression scores
    """
    # Map to manifold if needed (static branch - JIT friendly)
    if self.input_space == "tangent":
        x = jax.vmap(self.manifold.expmap_0, in_axes=(0, None), out_axes=0)(x, c)

    # Compute multinomial linear regression
    res = self.manifold.compute_mlr(
        x,
        self.kernel[...],
        self.bias[...],
        c,
        self.clamping_factor,
        self.smoothing_factor,
    )

    return res

hyperbolix.nn_layers.HypRegressionHyperboloidBusemann

HypRegressionHyperboloidBusemann(
    manifold_module: Hyperboloid,
    in_dim: int,
    out_dim: int,
    *,
    rngs: Rngs,
    input_space: str = "manifold",
    param_dtype: DTypeLike = jnp.float32,
)

Bases: Module

Busemann MLR classification head (Hyperboloid / Lorentz model).

Computes per-class logits u_k(x) = -alpha_k·B^{v_k}(x) + b_k from the Lorentz Busemann function (point-to-horosphere distance), returning Euclidean logits — not manifold points.

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	Number of classes.	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'
param_dtype	DTypeLike	Storage dtype of the trainable parameters (default: jnp.float32). Compute precision of manifold operations is set by manifold.dtype.	float32

Notes

Parameters use the weight-normalization split (kernel directions, log_scale log-magnitudes alpha = exp(log_scale), bias); see :func:busemann_core.busemann_score.

References

Chen, Schölkopf, and Sebe. "Hyperbolic Busemann Neural Networks." 2026, Sec. 3.

Source code in hyperbolix/nn_layers/busemann_regression.py

def __init__(
    self,
    manifold_module: Hyperboloid,
    in_dim: int,
    out_dim: int,
    *,
    rngs: nnx.Rngs,
    input_space: str = "manifold",
    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_hyperboloid_manifold(manifold_module, required_methods=("expmap_0", "busemann"))
    self.manifold = manifold_module
    self.in_dim = in_dim
    self.out_dim = out_dim
    self.input_space = input_space

    # Direction dim is spatial (ambient - 1) for the Lorentz model.
    in_spatial = in_dim - 1
    self.kernel, self.log_scale, self.bias = init_weight_norm_params(
        rngs, out_dim, in_spatial, std=in_spatial**-0.5, param_dtype=param_dtype
    )

__call__

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

Forward pass returning Euclidean Busemann logits, shape (batch, out_dim).

Source code in hyperbolix/nn_layers/busemann_regression.py

def __call__(
    self,
    x: Float[Array, "batch in_dim"],
    c: float = 1.0,
) -> Float[Array, "batch out_dim"]:
    """Forward pass returning Euclidean Busemann logits, shape (batch, out_dim)."""
    return busemann_score(self.manifold, x, self.kernel[...], self.log_scale[...], self.bias[...], c, self.input_space)

hyperbolix.nn_layers.FGGLorentzMLR

FGGLorentzMLR(
    in_features: int,
    num_classes: int,
    *,
    rngs: Rngs,
    reset_params: str = "mlr",
    init_bias: float = 0.5,
    eps: float = 1e-07,
    param_dtype: DTypeLike = jnp.float32,
)

Bases: Module

Fast and Geometrically Grounded Lorentz multinomial logistic regression.

Outputs Euclidean logits (signed scaled distances to hyperplanes) using the FGG spacelike V construction. Unlike HypRegressionHyperboloid, this layer uses the distance-to-hyperplane formulation matching the reference fc_mlr (signed_dist2hyperplanes_scaled_angle).

Forward pass (matching reference fc_mlr): 1. Build V_mink from (z, a) 2. mink = x @ V_mink (Minkowski inner products) 3. logits = asinh(sqrt(c) * mink) / sqrt(c) (signed scaled distances)

Parameters:

Name	Type	Description	Default
in_features	int	Input ambient dimension (D_in + 1), including time component.	required
num_classes	int	Number of output classes.	required
rngs	Rngs	Random number generators for parameter initialization.	required
reset_params	str	Weight initialization scheme for hyperplane normals: "mlr" (normal, std=sqrt(5/I)) or "default" (uniform) (default: "mlr").	'mlr'
init_bias	float	Initial value for bias entries (default: 0.5).	0.5
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), Eq. 23.

Source code in hyperbolix/nn_layers/hyperboloid_regression.py

def __init__(
    self,
    in_features: int,
    num_classes: int,
    *,
    rngs: nnx.Rngs,
    reset_params: str = "mlr",
    init_bias: float = 0.5,
    eps: float = 1e-7,
    param_dtype: DTypeLike = jnp.float32,
):
    if reset_params not in ("default", "mlr"):
        raise ValueError(f"reset_params must be 'default' or 'mlr', got '{reset_params}'")

    in_spatial = in_features - 1  # I
    self.in_features = in_features
    self.num_classes = num_classes
    self.eps = eps

    # Hyperplane normals (spatial) and bias offsets
    # Reference computes std from ambient dimension (in_features)
    key = rngs.params()
    if reset_params == "mlr":
        std = jnp.sqrt(5.0 / in_features)
        self.kernel = nnx.Param(jax.random.normal(key, (in_spatial, num_classes), dtype=param_dtype) * std)
    else:  # default
        stdv = 1.0 / jnp.sqrt(jnp.array(in_features, dtype=jnp.float32))
        self.kernel = nnx.Param(
            jax.random.uniform(key, (in_spatial, num_classes), dtype=param_dtype, minval=-stdv, maxval=stdv)
        )
    self.bias = nnx.Param(jnp.full((num_classes,), init_bias, dtype=param_dtype))

__call__

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

Forward pass returning Euclidean logits.

Parameters:

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

Returns:

Name	Type	Description
logits_BK	(Array, shape(B, K))	Euclidean logits (signed scaled distances to hyperplanes).

Source code in hyperbolix/nn_layers/hyperboloid_regression.py

def __call__(
    self,
    x_BAi: Float[Array, "batch in_features"],
    c: float = 1.0,
) -> Float[Array, "batch num_classes"]:
    """Forward pass returning Euclidean logits.

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

    Returns
    -------
    logits_BK : Array, shape (B, K)
        Euclidean logits (signed scaled distances to hyperplanes).
    """
    # 1. Build V_mink from (z, a)
    V_AiK = build_spacelike_V(self.kernel[...], self.bias[...], c, self.eps)  # (Ai, K)
    # Cast V to match input dtype (avoids float32/float64 scatter warnings)
    V_AiK = V_AiK.astype(x_BAi.dtype)

    # 2. Minkowski inner products
    mink_BK = x_BAi @ V_AiK  # (B, K)

    # 3. Signed scaled distances (matching reference fc_mlr: no norm scaling)
    sqrt_c = jnp.sqrt(c)
    logits_BK = jnp.asinh(sqrt_c * mink_BK) / sqrt_c  # (B, K)

    return logits_BK

Proper Velocity

HypRegressionPV (Chen et al. 2026, Thm 5.2 / Eq. 19) returns the Euclidean signed margin to each PV hyperplane — feed to a standard softmax cross-entropy loss.

hyperbolix.nn_layers.HypRegressionPV

HypRegressionPV(
    manifold_module: ProperVelocity,
    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 multinomial logistic regression layer (Proper Velocity model).

Implements Chen et al. 2026, Thm 5.2 / Eq. 19, delegating to ProperVelocity.compute_mlr. The output is the Euclidean signed margin to each PV hyperplane and is intended to be consumed by a standard softmax-cross-entropy classification loss.

Computation steps: 0) Optionally lift tangent-space input via expmap_0. 1) Compute the PV MLR scores via manifold.compute_mlr.

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 (number of classes).	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 MLR output (default: 1.0).	1.0
smoothing_factor	float	Smoothing factor for the MLR 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: Configuration parameters (input_space, clamping_factor, smoothing_factor) 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_regression.py

def __init__(
    self,
    manifold_module: ProperVelocity,
    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}'")

    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.clamping_factor = clamping_factor
    self.smoothing_factor = smoothing_factor

    # Kernel init: small normal with std = 1e-2 (matches paper reference
    # PVManifoldMLR.reset_parameters in the Chen et al. repo).
    self.kernel = nnx.Param(jax.random.normal(rngs.params(), (out_dim, in_dim), dtype=param_dtype) * 1e-2)
    # Bias init: uniform U(-1e-3, 1e-3) (matches paper reference).
    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 regression layer.

Parameters:

Name	Type	Description	Default
x	Array of shape (batch, in_dim)	Input tensor — PV manifold points or tangent-space vectors.	required
c	float	Manifold curvature (default: 1.0).	1.0

Returns:

Name	Type	Description
res	Array of shape (batch, out_dim)	MLR logits (Euclidean, suitable for cross-entropy).

Source code in hyperbolix/nn_layers/pv_regression.py

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

    Parameters
    ----------
    x : Array of shape (batch, in_dim)
        Input tensor — PV manifold points or tangent-space vectors.
    c : float
        Manifold curvature (default: 1.0).

    Returns
    -------
    res : Array of shape (batch, out_dim)
        MLR logits (Euclidean, suitable for cross-entropy).
    """
    if self.input_space == "tangent":
        x = jax.vmap(self.manifold.expmap_0, in_axes=(0, None), out_axes=0)(x, c)

    return self.manifold.compute_mlr(
        x,
        self.kernel[...],
        self.bias[...],
        c,
        self.clamping_factor,
        self.smoothing_factor,
    )

Example

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

poincare = Poincare()
head = HypRegressionPoincare(manifold_module=poincare, in_dim=32, out_dim=10, rngs=nnx.Rngs(0))

x = jax.random.normal(jax.random.PRNGKey(1), (64, 32)) * 0.3
x_proj = jax.vmap(poincare.proj, in_axes=(0, None))(x, 1.0)
logits = head(x_proj, c=1.0)            # (64, 10)
probs = jax.nn.softmax(logits, axis=-1)