Convolutional Layers

2D hyperbolic convolutions and pooling for the Hyperboloid, Poincaré, and Proper Velocity models. For which conv to pick see the NN Layers guide.

Dimensional growth (Hyperboloid HCat convolutions)

HypConv2DHyperboloid / HypConv2DHyperboloidFHNN / FGGConv2D increase dimensionality via the HCat operation: input d+1 → output (d×N)+1 where N = kernel_height × kernel_width. A 3×3 kernel turns a 3D input into 28D. Use small kernels or add dimensionality reduction. Poincaré, PV, and the HRC-based LorentzConv2D preserve dimension instead.

Hyperboloid

HypConv2DHyperboloid (the robust default) uses HCat patch extraction + Lorentz FC (Bdeir et al. 2023). HypConv2DHyperboloidILNN is the intrinsic Lorentz conv (Shi et al. 2026): LogCat (log-radius-preserving concatenation) + PLFC, with origin padding. LorentzConv2D is the dimension-preserving HRC-based variant (faster, lower accuracy — legacy/benchmarking).

hyperbolix.nn_layers.HypConv2DHyperboloid

HypConv2DHyperboloid(
    manifold_module: Hyperboloid,
    in_channels: int,
    out_channels: int,
    kernel_size: int | tuple[int, int],
    *,
    rngs: Rngs,
    stride: int | tuple[int, int] = 1,
    padding: str = "SAME",
    input_space: str = "manifold",
    reset_params: str = "default",
    param_dtype: DTypeLike = jnp.float32,
)

Bases: Module

Hyperbolic 2D Convolutional Layer for Hyperboloid model.

This layer implements fully hyperbolic convolution as described in "Fully Hyperbolic Convolutional Neural Networks for Computer Vision".

Computation steps: 1) Extract receptive field (kernel_size x kernel_size) of hyperbolic points 2) Apply HCat (Lorentz direct concatenation) to combine receptive field points 3) Pass through hyperbolic linear layer (LFC)

Parameters:

Name	Type	Description	Default
manifold_module	object	Class-based Hyperboloid manifold instance	required
in_channels	int	Number of input channels	required
out_channels	int	Number of output channels	required
kernel_size	int or tuple[int, int]	Size of the convolutional kernel	required
rngs	Rngs	Random number generators for parameter initialization	required
stride	int or tuple[int, int]	Stride of the convolution (default: 1)	1
padding	str	Padding mode, either 'SAME' or 'VALID' (default: 'SAME')	'SAME'
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'
reset_params	str	Kernel initialization scheme (default: "default"). "default": uniform U(-0.02, 0.02) (unchanged FHCNN-style init). "fan_in": normal, std = sqrt(1 / hcat_out_ambient_dim) — a norm-preserving fan-in init (analogous to _fgg_weight_init's "lorentz_kaiming" branch) validated for BN-free encoders.	'default'
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 (padding, input_space) 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_conv.py

def __init__(
    self,
    manifold_module: Hyperboloid,
    in_channels: int,
    out_channels: int,
    kernel_size: int | tuple[int, int],
    *,
    rngs: nnx.Rngs,
    stride: int | tuple[int, int] = 1,
    padding: str = "SAME",
    input_space: str = "manifold",
    reset_params: str = "default",
    param_dtype: DTypeLike = jnp.float32,
):
    if padding not in ["SAME", "VALID"]:
        raise ValueError(f"padding must be either 'SAME' or 'VALID', got '{padding}'")
    if input_space not in ["tangent", "manifold"]:
        raise ValueError(f"input_space must be either 'tangent' or 'manifold', got '{input_space}'")
    if reset_params not in ("default", "fan_in"):
        raise ValueError(f"reset_params must be 'default' or 'fan_in', got '{reset_params}'")

    # Static configuration
    validate_hyperboloid_manifold(manifold_module, required_methods=("expmap_0", "hcat"))
    self.manifold = manifold_module
    self.in_channels = in_channels
    self.out_channels = out_channels
    self.input_space = input_space
    self.padding = padding

    self.kernel_size = as_pair(kernel_size)
    self.stride = as_pair(stride)

    # HCat output ambient dim for the linear layer
    hcat_out_ambient_dim = hcat_ambient_dim(in_channels, self.kernel_size)

    # Trainable parameters — owned directly for flat parameter paths
    if reset_params == "default":
        bound = 0.02
        kernel_init = jax.random.uniform(
            rngs.params(), (out_channels, hcat_out_ambient_dim), dtype=param_dtype, minval=-bound, maxval=bound
        )
    else:  # "fan_in"
        std = jnp.sqrt(1.0 / hcat_out_ambient_dim)
        kernel_init = jax.random.normal(rngs.params(), (out_channels, hcat_out_ambient_dim), dtype=param_dtype) * std
    self.kernel = nnx.Param(kernel_init)
    self.bias = nnx.Param(jnp.zeros((1, out_channels), dtype=param_dtype))
    self.scale = 2.3  # not learnable (matches FHCNN default)

__call__

__call__(
    x: Float[Array, "batch height width in_channels"],
    c: float = 1.0,
) -> Float[
    Array, "batch out_height out_width out_channels"
]

Forward pass through the hyperbolic convolutional layer.

Parameters:

Name	Type	Description	Default
x	Array of shape (batch, height, width, in_channels)	Input feature map where each pixel is a point on the Hyperboloid manifold	required
c	float	Manifold curvature (default: 1.0)	1.0

Returns:

Name	Type	Description
out	Array of shape (batch, out_height, out_width, out_channels)	Output feature map on the Hyperboloid manifold

Source code in hyperbolix/nn_layers/hyperboloid_conv.py

def __call__(
    self,
    x: Float[Array, "batch height width in_channels"],
    c: float = 1.0,
) -> Float[Array, "batch out_height out_width out_channels"]:
    """
    Forward pass through the hyperbolic convolutional layer.

    Parameters
    ----------
    x : Array of shape (batch, height, width, in_channels)
        Input feature map where each pixel is a point on the Hyperboloid manifold
    c : float
        Manifold curvature (default: 1.0)

    Returns
    -------
    out : Array of shape (batch, out_height, out_width, out_channels)
        Output feature map on the Hyperboloid manifold
    """
    # Map to manifold if needed (static branch - JIT friendly)
    x = _map_input_to_manifold(x, self.manifold, self.input_space, c)

    # Extract patches: (B, H, W, kh, kw, C)
    patches_BHWkhkwC = extract_patches(x, self.kernel_size, self.stride, self.padding, pad_mode="edge")
    batch, out_h, out_w, kh, kw, in_c = patches_BHWkhkwC.shape

    # Flatten batch+spatial for parallel processing: (B*H*W, K, C)
    patches_flat_NKC = patches_BHWkhkwC.reshape(-1, kh * kw, in_c)

    # HCat: (K, C) -> (hcat_dim,) per patch
    hcat_out_NA = jax.vmap(self.manifold.hcat, in_axes=(0, None))(patches_flat_NKC, c)  # (B*H*W, hcat_dim)

    # Linear: (hcat_dim,) -> (out_channels,)
    linear_out_NC = _fhcnn_forward(
        hcat_out_NA,
        self.kernel[...],
        self.bias[...],
        self.manifold,
        c,
        "manifold",
        None,
        False,
        self.scale,
        1e-5,
    )  # (B*H*W, out_channels)

    # Reshape back to spatial
    output_BHWC = linear_out_NC.reshape(batch, out_h, out_w, self.out_channels)

    return output_BHWC

hyperbolix.nn_layers.HypConv2DHyperboloidFHNN

HypConv2DHyperboloidFHNN(
    manifold_module: Hyperboloid,
    in_channels: int,
    out_channels: int,
    kernel_size: int | tuple[int, int],
    *,
    rngs: Rngs,
    stride: int | tuple[int, int] = 1,
    padding: str = "SAME",
    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 2D convolutional layer (Chen et al. 2021).

Uses HCat (Lorentz direct concatenation) to combine receptive field points, then applies the FHNN linear transform (time-primary sigmoid parameterization) for channel mixing.

Computation steps: 1) Map input to manifold via expmap_0 if input_space="tangent" 2) Extract receptive field (kernel_size x kernel_size) of hyperbolic points 3) Apply HCat (Lorentz direct concatenation) to combine receptive field points 4) Pass through FHNN linear (sigmoid time + spatial rescaling)

Parameters:

Name	Type	Description	Default
manifold_module	Hyperboloid	Class-based Hyperboloid manifold instance	required
in_channels	int	Number of input channels (ambient dimension, including time component)	required
out_channels	int	Number of output channels (ambient dimension, including time component)	required
kernel_size	int or tuple[int, int]	Size of the convolutional kernel	required
rngs	Rngs	Random number generators for parameter initialization	required
stride	int or tuple[int, int]	Stride of the convolution (default: 1)	1
padding	str	Padding mode, either 'SAME' or 'VALID' (default: 'SAME')	'SAME'
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).	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 (padding, input_space, activation) are treated as static and baked into the compiled function.

See Also

HypConv2DHyperboloid : Equivalent convolution using FHCNN linear instead of FHNN. HypLinearHyperboloidFHNN : The underlying FHNN linear layer.

References

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

Source code in hyperbolix/nn_layers/hyperboloid_conv.py

def __init__(
    self,
    manifold_module: Hyperboloid,
    in_channels: int,
    out_channels: int,
    kernel_size: int | tuple[int, int],
    *,
    rngs: nnx.Rngs,
    stride: int | tuple[int, int] = 1,
    padding: str = "SAME",
    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 padding not in ["SAME", "VALID"]:
        raise ValueError(f"padding must be either 'SAME' or 'VALID', got '{padding}'")
    if input_space not in ["tangent", "manifold"]:
        raise ValueError(f"input_space must be either 'tangent' or 'manifold', got '{input_space}'")

    # Static configuration
    validate_hyperboloid_manifold(manifold_module, required_methods=("expmap_0", "hcat"))
    self.manifold = manifold_module
    self.in_channels = in_channels
    self.out_channels = out_channels
    self.input_space = input_space
    self.padding = padding
    self.eps = eps
    self.activation = activation

    self.kernel_size = as_pair(kernel_size)
    self.stride = as_pair(stride)

    # HCat output ambient dim
    hcat_out_ambient_dim = hcat_ambient_dim(in_channels, self.kernel_size)

    # 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_channels, hcat_out_ambient_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_channels), 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 height width in_channels"],
    c: float = 1.0,
    deterministic: bool = True,
) -> Float[
    Array, "batch out_height out_width out_channels"
]

Forward pass through the FHNN hyperbolic convolutional layer.

Parameters:

Name	Type	Description	Default
x	Array of shape (batch, height, width, in_channels)	Input feature map where each pixel is a point on the Hyperboloid manifold	required
c	float	Manifold curvature (default: 1.0)	1.0
deterministic	bool	If True, dropout is disabled (default: True).	True

Returns:

Name	Type	Description
out	Array of shape (batch, out_height, out_width, out_channels)	Output feature map on the Hyperboloid manifold

Source code in hyperbolix/nn_layers/hyperboloid_conv.py

def __call__(
    self,
    x: Float[Array, "batch height width in_channels"],
    c: float = 1.0,
    deterministic: bool = True,
) -> Float[Array, "batch out_height out_width out_channels"]:
    """Forward pass through the FHNN hyperbolic convolutional layer.

    Parameters
    ----------
    x : Array of shape (batch, height, width, in_channels)
        Input feature map where each pixel is a point on the Hyperboloid manifold
    c : float
        Manifold curvature (default: 1.0)
    deterministic : bool
        If True, dropout is disabled (default: True).

    Returns
    -------
    out : Array of shape (batch, out_height, out_width, out_channels)
        Output feature map on the Hyperboloid manifold
    """
    # Map to manifold if needed (static branch - JIT friendly)
    x = _map_input_to_manifold(x, self.manifold, self.input_space, c)

    # Extract patches: (B, H, W, kh, kw, C)
    patches_BHWkhkwC = extract_patches(x, self.kernel_size, self.stride, self.padding, pad_mode="edge")
    batch, out_h, out_w, kh, kw, in_c = patches_BHWkhkwC.shape

    # Flatten batch+spatial for parallel processing: (B*H*W, K, C)
    patches_flat_NKC = patches_BHWkhkwC.reshape(-1, kh * kw, in_c)

    # HCat: (K, C) -> (hcat_dim,) per patch
    hcat_out_NA = jax.vmap(self.manifold.hcat, in_axes=(0, None))(patches_flat_NKC, c)  # (B*H*W, hcat_dim)

    # 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

    # FHNN linear: (hcat_dim,) -> (out_channels,)
    linear_out_NC = _fhnn_forward(
        hcat_out_NA,
        self.kernel[...],
        self.bias[...],
        self.manifold,
        c,
        "manifold",  # HCat output is already on manifold
        self.activation,
        dropout_fn,
        self.scale[...],
        self.eps,
    )  # (B*H*W, out_channels)

    # Reshape back to spatial
    return linear_out_NC.reshape(batch, out_h, out_w, self.out_channels)

hyperbolix.nn_layers.HypConv2DHyperboloidILNN

HypConv2DHyperboloidILNN(
    manifold_module: Hyperboloid,
    in_channels: int,
    out_channels: int,
    kernel_size: int | tuple[int, int],
    *,
    rngs: Rngs,
    stride: int | tuple[int, int] = 1,
    padding: str = "SAME",
    pad_mode: str = "origin",
    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

Intrinsic Lorentz Neural Network (ILNN) 2D convolutional layer (Hyperboloid model).

Implements the Lorentz convolution of Shi et al. 2026, Eq. (11): y = PLFC(LogCat({x_patch})). Receptive-field points are combined with the log-radius-preserving concatenation (LogCat) and mixed across channels by the PLFC linear transform (MLR scores -> sinh diffeomorphism -> time reconstruction). This extends the HNN++ linearized-kernel convolution (Shimizu et al. 2020) by replacing the Euclidean FC and naive concatenation with their intrinsic Lorentz counterparts. Formerly named HypConv2DHyperboloidPP.

Computation steps: 1) Extract receptive field (kernel_size x kernel_size) of hyperbolic points, padding with the manifold origin if needed 2) Apply LogCat (log-radius-preserving concatenation) to combine receptive field points 3) Pass through PLFC linear (MLR scores -> guard -> sinh -> time reconstruction) 4) Optionally add an intrinsic gyro-bias y <- y (+) exp_0([0, b])

Parameters:

Name	Type	Description	Default
manifold_module	Hyperboloid	Class-based Hyperboloid manifold instance	required
in_channels	int	Number of input channels (ambient dimension, including time component)	required
out_channels	int	Number of output channels (ambient dimension, including time component)	required
kernel_size	int or tuple[int, int]	Size of the convolutional kernel	required
rngs	Rngs	Random number generators for parameter initialization	required
stride	int or tuple[int, int]	Stride of the convolution (default: 1)	1
padding	str	Padding mode, either 'SAME' or 'VALID' (default: 'SAME')	'SAME'
pad_mode	str	How to fill padding pixels for 'SAME' padding: "origin" fills with the manifold origin (sqrt(1/c), 0, ..., 0) (matching the Shi et al. 2026 reference, which clamps zero-padded times up to sqrt(1/c)), "edge" replicates border values (default: "origin").	'origin'
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 smooth-clamped 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; the reference conv defers to the PLFC init). Use 1.0 to recover the previous HNN++-style init (Shimizu et al. 2020); note that large kernels push the pre-guard scores into the v_max saturation regime. Low end: in BN-free residual stacks (no GyroBN between blocks), every block becomes a strict contraction below std ~= 0.05 (probe-measured); std=0.1 is healthy. The 0.02 default is tuned for the GyroBN'd PLFC recipe (this reference pairs the conv with a following BatchNorm) and should NOT be treated as a safe default for unnormalized residual stacks — raise kernel_init_std (e.g. to 0.1) when omitting normalization.	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 (padding, pad_mode, input_space, clamping_factor, smoothing_factor, v_max) are treated as static and baked into the compiled function.

See Also

HypLinearHyperboloidPLFC : The underlying PLFC linear layer. hyperbolix.manifolds.hyperboloid.Hyperboloid.log_radius_concat : The LogCat operation.

References

Xianglong Shi, Ziheng Chen, Yunhan Jiang, and Nicu Sebe. "Intrinsic Lorentz Neural Network." ICLR 2026, arXiv:2602.23981 (Sec. 4.3: Lorentz convolution, log-radius concatenation). Shimizu Ryohei, Yusuke Mukuta, and Tatsuya Harada. "Hyperbolic neural networks++." arXiv preprint arXiv:2006.08210 (2020).

Source code in hyperbolix/nn_layers/hyperboloid_conv.py

def __init__(
    self,
    manifold_module: Hyperboloid,
    in_channels: int,
    out_channels: int,
    kernel_size: int | tuple[int, int],
    *,
    rngs: nnx.Rngs,
    stride: int | tuple[int, int] = 1,
    padding: str = "SAME",
    pad_mode: str = "origin",
    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 padding not in ["SAME", "VALID"]:
        raise ValueError(f"padding must be either 'SAME' or 'VALID', got '{padding}'")
    if pad_mode not in ("origin", "edge"):
        raise ValueError(f"pad_mode must be 'origin' or 'edge', got '{pad_mode}'")
    if input_space not in ["tangent", "manifold"]:
        raise ValueError(f"input_space must be either 'tangent' or 'manifold', got '{input_space}'")

    # Static configuration
    required_methods = ("expmap_0", "compute_mlr", "log_radius_concat")
    if use_gyro_bias:
        required_methods = ("expmap_0", "compute_mlr", "log_radius_concat", "embed_spatial_0", "addition")
    validate_hyperboloid_manifold(manifold_module, required_methods=required_methods)
    self.manifold = manifold_module
    self.in_channels = in_channels
    self.out_channels = out_channels
    self.input_space = input_space
    self.padding = padding
    self.pad_mode = pad_mode
    self.clamping_factor = clamping_factor
    self.smoothing_factor = smoothing_factor
    _assert_v_max_safe(v_max)
    self.v_max = v_max

    self.kernel_size = as_pair(kernel_size)
    self.stride = as_pair(stride)

    # LogCat output ambient dim (same as HCat)
    logcat_out_ambient_dim = hcat_ambient_dim(in_channels, self.kernel_size)

    # Trainable parameters — small normal init (Shi et al. 2026 PLFC reference;
    # the reference conv defers to the PLFC reset_parameters)
    out_spatial = out_channels - 1
    logcat_spatial = logcat_out_ambient_dim - 1
    kernel_init = kernel_init_std * jax.random.normal(rngs.params(), (out_spatial, logcat_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 height width in_channels"],
    c: float = 1.0,
) -> Float[
    Array, "batch out_height out_width out_channels"
]

Forward pass through the HNN++ hyperboloid convolutional layer.

Parameters:

Name	Type	Description	Default
x	Array of shape (batch, height, width, in_channels)	Input feature map where each pixel is a point on the Hyperboloid manifold	required
c	float	Manifold curvature (default: 1.0)	1.0

Returns:

Name	Type	Description
out	Array of shape (batch, out_height, out_width, out_channels)	Output feature map on the Hyperboloid manifold

Source code in hyperbolix/nn_layers/hyperboloid_conv.py

def __call__(
    self,
    x: Float[Array, "batch height width in_channels"],
    c: float = 1.0,
) -> Float[Array, "batch out_height out_width out_channels"]:
    """
    Forward pass through the HNN++ hyperboloid convolutional layer.

    Parameters
    ----------
    x : Array of shape (batch, height, width, in_channels)
        Input feature map where each pixel is a point on the Hyperboloid manifold
    c : float
        Manifold curvature (default: 1.0)

    Returns
    -------
    out : Array of shape (batch, out_height, out_width, out_channels)
        Output feature map on the Hyperboloid manifold
    """
    # Map to manifold if needed (static branch - JIT friendly)
    x = _map_input_to_manifold(x, self.manifold, self.input_space, c)

    # Extract patches: (B, H', W', kh, kw, C)
    patches_BHWkhkwC = extract_patches(x, self.kernel_size, self.stride, self.padding, self.pad_mode, c)
    batch, out_h, out_w, kh, kw, in_c = patches_BHWkhkwC.shape

    # Flatten batch+spatial for parallel processing: (B*H'*W', K, C)
    patches_flat_NKC = patches_BHWkhkwC.reshape(-1, kh * kw, in_c)

    # LogCat: (K, C) -> (logcat_dim,) per patch — log-radius-preserving concatenation (Shi et al. 2026, Eq. 11)
    logcat_out_NA = jax.vmap(self.manifold.log_radius_concat, in_axes=(0, None))(patches_flat_NKC, c)

    # PLFC linear: (logcat_dim,) -> (out_channels,)
    linear_out_NC = _hyperboloid_plfc_forward(
        logcat_out_NA,
        self.kernel[...],
        self.bias[...],
        self.gyro_bias[...] if self.gyro_bias is not None else None,
        self.manifold,
        c,
        "manifold",  # LogCat output is already on manifold
        self.clamping_factor,
        self.smoothing_factor,
        self.v_max,
    )  # (B*H'*W', out_channels)

    # Reshape back to spatial
    return linear_out_NC.reshape(batch, out_h, out_w, self.out_channels)

hyperbolix.nn_layers.FGGConv2D

FGGConv2D(
    manifold_module: Hyperboloid,
    in_channels: int,
    out_channels: int,
    kernel_size: int | tuple[int, int],
    *,
    rngs: Rngs,
    stride: int | tuple[int, int] = 1,
    padding: str = "SAME",
    pad_mode: str = "origin",
    activation: Callable | 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 2D convolutional layer.

Uses HCat (Lorentz direct concatenation) to combine receptive field points, then applies FGGLinear for the channel mixing. This matches the reference implementation pattern from Klis et al. 2026.

Computation steps: 1) Extract receptive field patches, pad with manifold origin if needed 2) Apply HCat (Lorentz direct concatenation) to combine patch points 3) Pass through FGGLinear for channel transformation

Parameters:

Name	Type	Description	Default
manifold_module	Hyperboloid	Class-based Hyperboloid manifold instance.	required
in_channels	int	Input ambient channels (D_in + 1), including time component.	required
out_channels	int	Output ambient channels (D_out + 1), including time component.	required
kernel_size	int or tuple[int, int]	Size of the convolutional kernel.	required
rngs	Rngs	Random number generators for parameter initialization.	required
stride	int or tuple[int, int]	Stride of the convolution (default: 1).	1
padding	str	Padding mode: "SAME" or "VALID" (default: "SAME").	'SAME'
pad_mode	str	How to fill padding pixels: "origin" fills with the manifold origin (sqrt(1/c), 0, ..., 0) (matching reference), "edge" replicates border values (default: "origin").	'origin'
activation	Callable or None	Euclidean activation for the FGGLinear (default: None).	None
reset_params	str	Weight init for FGGLinear: "eye", "xavier", "kaiming", "lorentz_kaiming", "fan_out", or "mlr" (default: "fan_out"). The "fan_out" default (std sqrt(1/out_spatial)) is norm-preserving, a deliberate deviation from the Klis et al. 2026 classification reference for unnormalized stacks (e.g. an RL backbone feeding a bounded projection). Pass reset_params="lorentz_kaiming", init_bias=0.5 to recover the previous init.	'fan_out'
use_weight_norm	bool	Weight normalization in FGGLinear (default: False). Note gain and the "fan_out" scale are renormalized away in this mode.	False
init_bias	float	Initial bias for FGGLinear (default: 0.0). A zero bias removes the ~sqrt(out) * init_bias time-row term; pair with "fan_out" for norm preservation.	0.0
gain	float	Multiplier on the random init (default: 1.0). With "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 manifold.dtype.	float32

References

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

Source code in hyperbolix/nn_layers/hyperboloid_conv.py

def __init__(
    self,
    manifold_module: Hyperboloid,
    in_channels: int,
    out_channels: int,
    kernel_size: int | tuple[int, int],
    *,
    rngs: nnx.Rngs,
    stride: int | tuple[int, int] = 1,
    padding: str = "SAME",
    pad_mode: str = "origin",
    activation: Callable | 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,
):
    if padding not in ("SAME", "VALID"):
        raise ValueError(f"padding must be either 'SAME' or 'VALID', got '{padding}'")
    if pad_mode not in ("origin", "edge"):
        raise ValueError(f"pad_mode must be 'origin' or 'edge', got '{pad_mode}'")

    validate_hyperboloid_manifold(manifold_module, required_methods=("hcat",))
    self.manifold = manifold_module
    self.in_channels = in_channels
    self.out_channels = out_channels
    self.padding = padding
    self.pad_mode = pad_mode
    self.eps = eps

    self.kernel_size = as_pair(kernel_size)
    self.stride = as_pair(stride)

    # HCat output ambient dim
    hcat_out_ambient = hcat_ambient_dim(in_channels, self.kernel_size)

    # Trainable parameters — owned directly for flat parameter paths
    in_spatial = hcat_out_ambient - 1  # I
    out_spatial = out_channels - 1  # O

    self.activation = activation
    self.use_weight_norm = use_weight_norm

    # 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, hcat_out_ambient, out_channels, reset_params, param_dtype, gain=gain
    )

    # Weight normalization: decompose kernel = softplus(kernel_scale) * kernel_dir / ||kernel_dir||
    if use_weight_norm:
        self.kernel_dir = nnx.Param(U_init)  # (I, O) direction
        g_init_val = jnp.sqrt(1.0 / (hcat_out_ambient + out_channels))
        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)

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

__call__

__call__(
    x: Float[Array, "batch height width in_channels"],
    c: float = 1.0,
) -> Float[
    Array, "batch out_height out_width out_channels"
]

Forward pass through the FGG convolutional layer.

Parameters:

Name	Type	Description	Default
x	(Array, shape(B, H, W, in_channels))	Input feature map on the hyperboloid.	required
c	float	Curvature parameter (default: 1.0).	1.0

Returns:

Name	Type	Description
out	Array, shape (B, H', W', out_channels)	Output feature map on the hyperboloid.

Source code in hyperbolix/nn_layers/hyperboloid_conv.py

def __call__(
    self,
    x: Float[Array, "batch height width in_channels"],
    c: float = 1.0,
) -> Float[Array, "batch out_height out_width out_channels"]:
    """Forward pass through the FGG convolutional layer.

    Parameters
    ----------
    x : Array, shape (B, H, W, in_channels)
        Input feature map on the hyperboloid.
    c : float, optional
        Curvature parameter (default: 1.0).

    Returns
    -------
    out : Array, shape (B, H', W', out_channels)
        Output feature map on the hyperboloid.
    """
    # Extract patches: (B, H', W', kh, kw, C)
    patches = extract_patches(x, self.kernel_size, self.stride, self.padding, self.pad_mode, c)
    batch, out_h, out_w, kh, kw, in_c = patches.shape

    # Flatten batch+spatial: (B*H'*W', K, C) where K = kh*kw
    patches_flat_NKC = patches.reshape(-1, kh * kw, in_c)

    # HCat: (K, C) -> (hcat_dim,) per patch
    hcat_out_NA = jax.vmap(self.manifold.hcat, in_axes=(0, None))(patches_flat_NKC, c)

    # FGG forward: (hcat_dim,) -> (out_channels,)
    U_IO = _get_effective_kernel(
        getattr(self, "kernel", None),
        getattr(self, "kernel_dir", None),
        getattr(self, "kernel_scale", None),
        self.use_weight_norm,
        self.eps,
    )
    linear_out_NC = _fgg_linear_forward(hcat_out_NA, U_IO, self.bias[...], c, self.activation, self.eps)

    # Reshape back to spatial
    return linear_out_NC.reshape(batch, out_h, out_w, self.out_channels)

hyperbolix.nn_layers.LorentzConv2D

LorentzConv2D(
    in_channels: int,
    out_channels: int,
    kernel_size: int | tuple[int, int],
    *,
    rngs: Rngs,
    stride: int | tuple[int, int] = 1,
    padding: str = "SAME",
    param_dtype: DTypeLike = jnp.float32,
)

Bases: Module

Lorentz 2D Convolutional Layer using the Hyperbolic Layer (HL) approach.

This layer applies convolution to the space-like components of Lorentzian vectors and reconstructs the time-like component to maintain the manifold constraint. This is equivalent to an HRC (Hyperbolic Regularization Component) wrapper around a standard Conv2D.

Computation steps: 1) Extract space-like components x_s from input x = [x_t, x_s]^T 2) Apply Euclidean convolution: y_s = Conv2D(x_s) 3) Reconstruct time component: y_t = sqrt(||y_s||^2 + 1/c) 4) Return y = [y_t, y_s]^T

Parameters:

Name	Type	Description	Default
in_channels	int	Number of input channels (ambient dimension, including time component)	required
out_channels	int	Number of output channels (ambient dimension, including time component)	required
kernel_size	int or tuple[int, int]	Size of the convolutional kernel	required
rngs	Rngs	Random number generators for parameter initialization	required
stride	int or tuple[int, int]	Stride of the convolution (default: 1)	1
padding	str or int or tuple	Padding mode: 'SAME', 'VALID', or explicit padding (default: 'SAME')	'SAME'
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

Notes

This implementation follows the Hyperbolic Layer (HL) approach from "Fully Hyperbolic Convolutional Neural Networks for Computer Vision".

The layer operates only on space-like components, making it more computationally efficient than the HCat-based approach (HypConv2DHyperboloid), though it doesn't perform true hyperbolic convolution. Instead, it applies Euclidean operations to spatial components and reconstructs the time component.

See Also

hypformer.hrc : Core HRC function this layer is based on HypConv2DHyperboloid : Full hyperbolic convolution using HCat concatenation

References

He, Neil, Menglin Yang, and Rex Ying. "Lorentzian residual neural networks." Proceedings of the 31st ACM SIGKDD Conference on Knowledge Discovery and Data Mining V. 1. 2025.

Source code in hyperbolix/nn_layers/hyperboloid_conv.py

def __init__(
    self,
    in_channels: int,
    out_channels: int,
    kernel_size: int | tuple[int, int],
    *,
    rngs: nnx.Rngs,
    stride: int | tuple[int, int] = 1,
    padding: str = "SAME",
    param_dtype: DTypeLike = jnp.float32,
):
    self.in_channels = in_channels
    self.out_channels = out_channels

    # Create Euclidean conv layer for space components only
    # in_channels - 1: skip time component at index 0
    # out_channels - 1: time will be reconstructed from constraint
    self.conv = nnx.Conv(
        in_features=in_channels - 1,
        out_features=out_channels - 1,
        kernel_size=kernel_size,
        strides=stride,
        padding=padding,
        param_dtype=param_dtype,
        rngs=rngs,
    )

__call__

__call__(
    x: Float[Array, "batch height width in_channels"],
    c: float = 1.0,
) -> Float[
    Array, "batch out_height out_width out_channels"
]

Forward pass through the Lorentz convolutional layer.

This layer is a specific instance of the Hyperbolic Regularization Component (HRC) where the regularization function f_r is a 2D convolution. The HRC pattern: 1. Extracts space components 2. Applies Euclidean convolution 3. Reconstructs time component using Lorentz constraint

Parameters:

Name	Type	Description	Default
x	Array of shape (batch, height, width, in_channels)	Input feature map where x[..., 0] is time component and x[..., 1:] are space components on the Lorentz manifold	required
c	float	Manifold curvature parameter (default: 1.0)	1.0

Returns:

Name	Type	Description
out	Array of shape (batch, out_height, out_width, out_channels)	Output feature map on the Lorentz manifold

Notes

This implementation uses the HRC function from hypformer.py, demonstrating that LorentzConv2D (from LResNet) and HRC (from Hypformer) are mathematically equivalent approaches to adapting Euclidean operations for hyperbolic geometry.

Source code in hyperbolix/nn_layers/hyperboloid_conv.py

def __call__(
    self,
    x: Float[Array, "batch height width in_channels"],
    c: float = 1.0,
) -> Float[Array, "batch out_height out_width out_channels"]:
    """
    Forward pass through the Lorentz convolutional layer.

    This layer is a specific instance of the Hyperbolic Regularization Component (HRC)
    where the regularization function f_r is a 2D convolution. The HRC pattern:
    1. Extracts space components
    2. Applies Euclidean convolution
    3. Reconstructs time component using Lorentz constraint

    Parameters
    ----------
    x : Array of shape (batch, height, width, in_channels)
        Input feature map where x[..., 0] is time component and
        x[..., 1:] are space components on the Lorentz manifold
    c : float
        Manifold curvature parameter (default: 1.0)

    Returns
    -------
    out : Array of shape (batch, out_height, out_width, out_channels)
        Output feature map on the Lorentz manifold

    Notes
    -----
    This implementation uses the HRC function from hypformer.py, demonstrating that
    LorentzConv2D (from LResNet) and HRC (from Hypformer) are mathematically equivalent
    approaches to adapting Euclidean operations for hyperbolic geometry.
    """

    # Define convolution as the HRC regularization function f_r
    def conv_fn(x_space):
        return self.conv(x_space)

    # Apply HRC with curvature-preserving transformation (c_in = c_out = c)
    return hrc(x, conv_fn, c_in=c, c_out=c, eps=1e-8)

Poincaré

HypConv2DPoincare extracts patches, applies beta-concatenation (HNN++, Shimizu et al. 2020) over the receptive field, then a HypLinearPoincarePP layer. Dimension is preserved (K² × C_in → C_out); I/O is in tangent space by default.

hyperbolix.nn_layers.HypConv2DPoincare

HypConv2DPoincare(
    manifold_module: Poincare,
    in_channels: int,
    out_channels: int,
    kernel_size: int | tuple[int, int],
    *,
    rngs: Rngs,
    stride: int | tuple[int, int] = 1,
    padding: str = "SAME",
    input_space: str = "tangent",
    id_init: bool = True,
    clamping_factor: float = 1.0,
    smoothing_factor: float = 50.0,
    param_dtype: DTypeLike = jnp.float32,
)

Bases: Module

Hyperbolic 2D Convolutional Layer for Poincaré ball model.

This layer implements hyperbolic convolution following the Poincaré ResNet (van Spengler et al. 2023) approach: beta-scaling in tangent space, patch extraction, expmap to manifold, HNN++ fully-connected, logmap back to tangent space.

The layer operates in tangent space internally and returns tangent-space output, matching the reference implementation. This design avoids the numerically unstable logmap_0 round-trips that cause NaN gradients when points approach the Poincaré ball boundary.

Computation steps: 1) Map to tangent space if input is on manifold 2) Scale tangent vectors by beta function ratio (beta-concatenation scaling) 3) Extract patches via im2col (zero-padding in tangent space) 4) Map concatenated patch vectors to manifold via expmap_0 5) Apply HNN++ fully-connected layer 6) Map back to tangent space via logmap_0

Parameters:

Name	Type	Description	Default
manifold_module	Poincare	Class-based Poincaré manifold instance	required
in_channels	int	Number of input channels (Poincaré ball dimension per pixel)	required
out_channels	int	Number of output channels (Poincaré ball dimension per pixel)	required
kernel_size	int or tuple[int, int]	Size of the convolutional kernel	required
rngs	Rngs	Random number generators for parameter initialization	required
stride	int or tuple[int, int]	Stride of the convolution (default: 1)	1
padding	str	Padding mode, either 'SAME' or 'VALID' (default: 'SAME')	'SAME'
input_space	str	Type of the input tensor, either 'tangent' or 'manifold' (default: 'tangent'). Note: This is a static configuration - changing it after initialization requires recompilation.	'tangent'
id_init	bool	If True, use identity initialization (1/2 * I) for the linear sublayer weights, matching the reference Poincaré ResNet implementation (default: True). The 1/2 factor compensates for the factor of 2 inside the HNN++ distance formula.	True
clamping_factor	float	Clamping factor for the HNN++ linear layer output (default: 1.0)	1.0
smoothing_factor	float	Smoothing factor for the HNN++ linear layer 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

Output Space: This layer always returns tangent-space output (matching the reference). Between conv layers, use standard activations (e.g., jax.nn.relu) directly on the tangent-space features. Use expmap_0 to map back to the manifold only when needed (e.g., before the classification head).

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

Dimension math: - beta-scaling + patch extraction: (H, W, C_in) → (oh, ow, K^2 * C_in) - HNN++ linear: in_dim = K^2 * C_in, out_dim = C_out

References

Shimizu et al. "Hyperbolic neural networks++." arXiv:2006.08210 (2020). van Spengler et al. "Poincaré ResNet." ICML 2023.

Source code in hyperbolix/nn_layers/poincare_conv.py

def __init__(
    self,
    manifold_module: Poincare,
    in_channels: int,
    out_channels: int,
    kernel_size: int | tuple[int, int],
    *,
    rngs: nnx.Rngs,
    stride: int | tuple[int, int] = 1,
    padding: str = "SAME",
    input_space: str = "tangent",
    id_init: bool = True,
    clamping_factor: float = 1.0,
    smoothing_factor: float = 50.0,
    param_dtype: DTypeLike = jnp.float32,
):
    if padding not in ["SAME", "VALID"]:
        raise ValueError(f"padding must be either 'SAME' or 'VALID', got '{padding}'")
    if input_space not in ["tangent", "manifold"]:
        raise ValueError(f"input_space must be either 'tangent' or 'manifold', got '{input_space}'")

    # Static configuration
    validate_poincare_manifold(
        manifold_module,
        required_methods=("expmap_0", "logmap_0", "compute_mlr_pp"),
    )
    self.manifold = manifold_module
    self.in_channels = in_channels
    self.out_channels = out_channels
    self.input_space = input_space
    self.padding = padding

    self.kernel_size = as_pair(kernel_size)
    self.stride = as_pair(stride)

    # Precompute beta function ratio for tangent-space scaling
    # B(n/2, 1/2) / B(n_i/2, 1/2) where n = K^2 * C_in, n_i = C_in
    kernel_h, kernel_w = self.kernel_size
    K2 = kernel_h * kernel_w
    concat_dim = K2 * in_channels
    beta_n = jax.scipy.special.beta(concat_dim / 2.0, 0.5)
    beta_ni = jax.scipy.special.beta(in_channels / 2.0, 0.5)
    self.beta_scale = float(beta_n / beta_ni)

    self.clamping_factor = clamping_factor
    self.smoothing_factor = smoothing_factor

    # Trainable parameters — owned directly for flat parameter paths
    # Weight initialization: identity from reference (van Spengler et al. 2023)
    # or scaled normal std = 1/sqrt(fan_in)
    if id_init:
        # W = 1/2 * I(C_out, K^2*C_in)
        # The 1/2 factor compensates for the factor of 2 in the HNN++ distance formula.
        kernel_init = 0.5 * jnp.eye(out_channels, concat_dim, dtype=param_dtype)
    else:
        std = 1.0 / jnp.sqrt(concat_dim)
        kernel_init = jax.random.normal(rngs.params(), (out_channels, concat_dim), dtype=param_dtype) * std
    self.kernel = nnx.Param(kernel_init)
    self.bias = nnx.Param(jnp.zeros((out_channels, 1), dtype=param_dtype))

__call__

__call__(
    x: Float[Array, "batch height width in_channels"],
    c: float = 1.0,
) -> Float[
    Array, "batch out_height out_width out_channels"
]

Forward pass through the Poincaré convolutional layer.

Follows the reference computation flow: tangent-space beta-scaling, patch extraction, expmap_0, HNN++ FC, logmap_0.

Parameters:

Name	Type	Description	Default
x	Array of shape (batch, height, width, in_channels)	Input feature map. Can be tangent-space or manifold points depending on input_space setting.	required
c	float	Manifold curvature (default: 1.0)	1.0

Returns:

Name	Type	Description
out	Array of shape (batch, out_height, out_width, out_channels)	Output feature map in tangent space at the origin. Use standard activations (e.g., jax.nn.relu) between layers.

Source code in hyperbolix/nn_layers/poincare_conv.py

def __call__(
    self,
    x: Float[Array, "batch height width in_channels"],
    c: float = 1.0,
) -> Float[Array, "batch out_height out_width out_channels"]:
    """
    Forward pass through the Poincaré convolutional layer.

    Follows the reference computation flow: tangent-space beta-scaling,
    patch extraction, expmap_0, HNN++ FC, logmap_0.

    Parameters
    ----------
    x : Array of shape (batch, height, width, in_channels)
        Input feature map. Can be tangent-space or manifold points
        depending on input_space setting.
    c : float
        Manifold curvature (default: 1.0)

    Returns
    -------
    out : Array of shape (batch, out_height, out_width, out_channels)
        Output feature map in tangent space at the origin.
        Use standard activations (e.g., jax.nn.relu) between layers.
    """
    # Step 1: Map to tangent space if input is on manifold
    if self.input_space == "manifold":
        orig_shape = x.shape
        x_flat_NC = x.reshape(-1, x.shape[-1])  # (B*H*W, C_in)
        x_flat_NC = jax.vmap(self.manifold.logmap_0, in_axes=(0, None))(x_flat_NC, c)
        x = x_flat_NC.reshape(orig_shape)

    # Now x is in tangent space: (B, H, W, C_in)

    # Step 2: Scale tangent vectors by beta ratio (matching reference)
    x = x * self.beta_scale

    # Step 3: Extract patches in tangent space using zero-padding
    kernel_h, kernel_w = self.kernel_size
    stride_h, stride_w = self.stride

    patches_BHWKC = jax.lax.conv_general_dilated_patches(
        lhs=x,
        filter_shape=(kernel_h, kernel_w),
        window_strides=(stride_h, stride_w),
        padding=self.padding,
        dimension_numbers=("NHWC", "OIHW", "NHWC"),
    )  # (B, H, W, K²·C_in)

    batch, out_h, out_w, concat_dim = patches_BHWKC.shape

    # Step 4: Map concatenated patch vectors to manifold via expmap_0
    patches_flat_NKC = patches_BHWKC.reshape(-1, concat_dim)  # (N, K²·C_in) where N=B*H*W
    manifold_pts_NKC = jax.vmap(self.manifold.expmap_0, in_axes=(0, None))(
        patches_flat_NKC, c
    )  # (N, K²·C_in) on Poincaré ball

    # Step 5: HNN++ FC — (N, K²·C_in) on manifold → (N, C_out) on manifold
    fc_out_NC = _poincare_pp_forward(
        manifold_pts_NKC,
        self.kernel[...],
        self.bias[...],
        self.manifold,
        c,
        "manifold",
        self.clamping_factor,
        self.smoothing_factor,
    )

    # Step 6: Map back to tangent space
    tangent_out_NC = jax.vmap(self.manifold.logmap_0, in_axes=(0, None))(fc_out_NC, c)  # (N, C_out)

    # Reshape to spatial output
    output_BHWC = tangent_out_NC.reshape(batch, out_h, out_w, self.out_channels)

    return output_BHWC

Proper Velocity

HypConv2DPV (Chen et al. 2026, Sec 5.3). Because PV geometry is unconstrained \(\mathbb{R}^n\), patch concatenation coincides with Euclidean concatenation — no beta-scaling step and no dimension growth. Outputs live on the PV manifold, so Euclidean activations apply directly between conv layers.

hyperbolix.nn_layers.HypConv2DPV

HypConv2DPV(
    manifold_module: ProperVelocity,
    in_channels: int,
    out_channels: int,
    kernel_size: int | tuple[int, int],
    *,
    rngs: Rngs,
    stride: int | tuple[int, int] = 1,
    padding: str = "SAME",
    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 2D Convolutional Layer for the Proper Velocity model.

Implements Chen et al. 2026, Sec 5.3. Because PV's geometry is unconstrained (ℝⁿ), patch concatenation coincides with Euclidean concatenation — no beta-scaling step is required (unlike HypConv2DPoincare). The layer returns points on the PV manifold; between conv layers, activations can be applied directly in PV space (paper Sec 5.3 "Activation").

Computation steps: 1) Optionally lift tangent-space input to the PV manifold via expmap_0. 2) Extract patches with zero-padding via jax.lax.conv_general_dilated_patches. 3) Flatten batch x spatial dimensions. 4) Apply the PV FC forward (_pv_fc_forward, shared with HypLinearPV). 5) Reshape back to spatial output.

Parameters:

Name	Type	Description	Default
manifold_module	ProperVelocity	Class-based Proper Velocity manifold instance.	required
in_channels	int	Number of input channels (PV dimension per pixel).	required
out_channels	int	Number of output channels (PV dimension per pixel).	required
kernel_size	int or tuple[int, int]	Size of the convolutional kernel.	required
rngs	Rngs	Random number generators for parameter initialization.	required
stride	int or tuple[int, int]	Stride of the convolution (default: 1).	1
padding	str	Padding mode, either 'SAME' or 'VALID' (default: 'SAME').	'SAME'
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 inside the outer sinh (paper Eq. 23). Default None.	None
clamping_factor	float	Clamping factor for the PV MLR output (default: 1.0).	1.0
smoothing_factor	float	Smoothing factor for the PV 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 / (kernel_h * kernel_w * in_channels)) so stacked conv blocks preserve variance under ReLU (for small MLR arguments the PV output reduces to an Euclidean linear map per patch, so standard He analysis applies).	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

Output Space: Unlike HypConv2DPoincare (which returns tangent space), this layer returns points on the PV manifold. Between conv layers, apply activations directly on the PV features — no expmap_0/logmap_0 round-trips needed.

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

Dimension math: - patch extraction: (H, W, C_in) → (oh, ow, K²·C_in) - PV FC: in_dim = K²·C_in, out_dim = C_out

References

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

Source code in hyperbolix/nn_layers/pv_conv.py

def __init__(
    self,
    manifold_module: ProperVelocity,
    in_channels: int,
    out_channels: int,
    kernel_size: int | tuple[int, int],
    *,
    rngs: nnx.Rngs,
    stride: int | tuple[int, int] = 1,
    padding: str = "SAME",
    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 padding not in ["SAME", "VALID"]:
        raise ValueError(f"padding must be either 'SAME' or 'VALID', got '{padding}'")
    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_channels = in_channels
    self.out_channels = out_channels
    self.input_space = input_space
    self.padding = padding
    self.inner_activation = inner_activation
    self.clamping_factor = clamping_factor
    self.smoothing_factor = smoothing_factor

    self.kernel_size = as_pair(kernel_size)
    self.stride = as_pair(stride)

    kernel_h, kernel_w = self.kernel_size
    concat_dim = kernel_h * kernel_w * in_channels

    # Kernel init: He(fan_in) by default so deep conv stacks preserve
    # variance under ReLU. Bias init is uniform U(-1e-3, 1e-3), matching
    # the paper reference.
    std = (2.0 / concat_dim) ** 0.5 if kernel_init_std is None else kernel_init_std
    self.kernel = nnx.Param(jax.random.normal(rngs.params(), (out_channels, concat_dim), dtype=param_dtype) * std)
    self.bias = nnx.Param(
        jax.random.uniform(rngs.params(), (out_channels, 1), dtype=param_dtype, minval=-1e-3, maxval=1e-3)
    )

__call__

__call__(
    x: Float[Array, "batch height width in_channels"],
    c: float = 1.0,
) -> Float[
    Array, "batch out_height out_width out_channels"
]

Forward pass through the PV convolutional layer.

Parameters:

Name	Type	Description	Default
x	Array of shape (batch, height, width, in_channels)	Input feature map. PV manifold points or tangent-space vectors at the origin, depending on input_space.	required
c	float	Manifold curvature (default: 1.0).	1.0

Returns:

Name	Type	Description
out	Array of shape (batch, out_height, out_width, out_channels)	Output feature map on the PV manifold.

Source code in hyperbolix/nn_layers/pv_conv.py

def __call__(
    self,
    x: Float[Array, "batch height width in_channels"],
    c: float = 1.0,
) -> Float[Array, "batch out_height out_width out_channels"]:
    """
    Forward pass through the PV convolutional layer.

    Parameters
    ----------
    x : Array of shape (batch, height, width, in_channels)
        Input feature map. PV manifold points or tangent-space vectors at the
        origin, depending on ``input_space``.
    c : float
        Manifold curvature (default: 1.0).

    Returns
    -------
    out : Array of shape (batch, out_height, out_width, out_channels)
        Output feature map on the PV manifold.
    """
    # Step 1: Optionally lift tangent-space input to the PV manifold.
    if self.input_space == "tangent":
        orig_shape = x.shape
        x_flat_NC = x.reshape(-1, x.shape[-1])  # (B*H*W, C_in)
        x_flat_NC = jax.vmap(self.manifold.expmap_0, in_axes=(0, None))(x_flat_NC, c)
        x = x_flat_NC.reshape(orig_shape)

    # Step 2: Extract patches with zero-padding (raw concat — no beta scaling).
    kernel_h, kernel_w = self.kernel_size
    stride_h, stride_w = self.stride

    patches_BHWKC = jax.lax.conv_general_dilated_patches(
        lhs=x,
        filter_shape=(kernel_h, kernel_w),
        window_strides=(stride_h, stride_w),
        padding=self.padding,
        dimension_numbers=("NHWC", "OIHW", "NHWC"),
    )  # (B, H, W, K²·C_in)

    batch, out_h, out_w, concat_dim = patches_BHWKC.shape

    # Step 3: Flatten batch x spatial dims for the FC.
    patches_flat_NKC = patches_BHWKC.reshape(-1, concat_dim)  # (N, K²·C_in)

    # Step 4: PV FC — input is already on the PV manifold (no extra lift).
    fc_out_NC = _pv_fc_forward(
        patches_flat_NKC,
        self.kernel[...],
        self.bias[...],
        self.manifold,
        c,
        "manifold",
        self.clamping_factor,
        self.smoothing_factor,
        self.inner_activation,
    )  # (N, C_out) on PV manifold

    # Step 5: Reshape to spatial output.
    return fc_out_NC.reshape(batch, out_h, out_w, self.out_channels)

Pooling

hyperbolix.nn_layers.hyp_avg_pool2d

hyp_avg_pool2d(
    x: Float[Array, "... H W dim_plus_1"],
    c: float,
    eps: float = 1e-07,
) -> Float[Array, "... dim_plus_1"]

Global average pooling over 2D spatial dimensions on the hyperboloid.

Averages the spatial components over the height and width dimensions, then reconstructs the time component via the hyperboloid constraint. This is the HRC pattern with f_r = mean_over_spatial:

.. math::

\text{space} = \text{mean}_{H,W}(x[..., 1:])

x_0 = \sqrt{\|\text{space}\|^2 + 1/c}

\text{output} = [x_0, \text{space}]

The function expects the NHWC layout used throughout hyperbolix: (..., H, W, ambient_dim) where ambient_dim = spatial_dim + 1.

This operation is hyperboloid-only because it relies on the time/spatial decomposition x[..., 0] (time) and x[..., 1:] (spatial). Poincaré ball points do not have a time component. To pool Poincaré features, first convert to hyperboloid via :func:~hyperbolix.manifolds.isometry_mappings.poincare_to_hyperboloid, pool, then convert back.

Parameters:

Name	Type	Description	Default
x	(Array, shape(..., H, W, dim + 1))	Hyperboloid feature map with curvature c. The last axis is the ambient dimension (time + spatial). The two axes before it are the spatial height H and width W.	required
c	float	Curvature parameter (positive, c > 0).	required
eps	float	Numerical stability floor for the time reconstruction (default: 1e-7).	1e-07

Returns:

Type	Description
(Array, shape(..., dim + 1))	Pooled hyperboloid points with curvature c.

See Also

spatial_to_hyperboloid : Low-level curvature scaling + time reconstruction. hrc : General HRC wrapper for arbitrary Euclidean functions. lorentz_midpoint : Weighted Lorentzian midpoint (geometrically exact aggregation).

Examples:

>>> import jax
>>> import jax.numpy as jnp
>>> from hyperbolix.manifolds import Hyperboloid
>>> from hyperbolix.nn_layers.hyperboloid_core import hyp_avg_pool2d
>>>
>>> hyperboloid = Hyperboloid(dtype=jnp.float32)
>>> key = jax.random.PRNGKey(0)
>>> v = jax.random.normal(key, (4, 7, 7, 64), dtype=jnp.float32) * 0.1
>>> x = jax.vmap(jax.vmap(jax.vmap(
...     hyperboloid.expmap_0, in_axes=(0, None)
... ), in_axes=(0, None)), in_axes=(0, None))(v, 1.0)
>>> x.shape
(4, 7, 7, 65)
>>> y = hyp_avg_pool2d(x, c=1.0)
>>> y.shape
(4, 65)

Source code in hyperbolix/nn_layers/hyperboloid_core.py

def hyp_avg_pool2d(
    x: Float[Array, "... H W dim_plus_1"],
    c: float,
    eps: float = 1e-7,
) -> Float[Array, "... dim_plus_1"]:
    """Global average pooling over 2D spatial dimensions on the hyperboloid.

    Averages the spatial components over the height and width dimensions,
    then reconstructs the time component via the hyperboloid constraint.
    This is the HRC pattern with ``f_r = mean_over_spatial``:

    .. math::

        \\text{space} = \\text{mean}_{H,W}(x[..., 1:])

        x_0 = \\sqrt{\\|\\text{space}\\|^2 + 1/c}

        \\text{output} = [x_0, \\text{space}]

    The function expects the **NHWC** layout used throughout hyperbolix:
    ``(..., H, W, ambient_dim)`` where ``ambient_dim = spatial_dim + 1``.

    This operation is **hyperboloid-only** because it relies on the time/spatial
    decomposition ``x[..., 0]`` (time) and ``x[..., 1:]`` (spatial). Poincaré
    ball points do not have a time component. To pool Poincaré features, first
    convert to hyperboloid via
    :func:`~hyperbolix.manifolds.isometry_mappings.poincare_to_hyperboloid`,
    pool, then convert back.

    Parameters
    ----------
    x : Array, shape (..., H, W, dim+1)
        Hyperboloid feature map with curvature ``c``.  The last axis is
        the ambient dimension (time + spatial).  The two axes before it
        are the spatial height ``H`` and width ``W``.
    c : float
        Curvature parameter (positive, c > 0).
    eps : float, optional
        Numerical stability floor for the time reconstruction
        (default: 1e-7).

    Returns
    -------
    Array, shape (..., dim+1)
        Pooled hyperboloid points with curvature ``c``.

    See Also
    --------
    spatial_to_hyperboloid : Low-level curvature scaling + time reconstruction.
    hrc : General HRC wrapper for arbitrary Euclidean functions.
    lorentz_midpoint : Weighted Lorentzian midpoint (geometrically exact aggregation).

    Examples
    --------
    >>> import jax
    >>> import jax.numpy as jnp
    >>> from hyperbolix.manifolds import Hyperboloid
    >>> from hyperbolix.nn_layers.hyperboloid_core import hyp_avg_pool2d
    >>>
    >>> hyperboloid = Hyperboloid(dtype=jnp.float32)
    >>> key = jax.random.PRNGKey(0)
    >>> v = jax.random.normal(key, (4, 7, 7, 64), dtype=jnp.float32) * 0.1
    >>> x = jax.vmap(jax.vmap(jax.vmap(
    ...     hyperboloid.expmap_0, in_axes=(0, None)
    ... ), in_axes=(0, None)), in_axes=(0, None))(v, 1.0)
    >>> x.shape
    (4, 7, 7, 65)
    >>> y = hyp_avg_pool2d(x, c=1.0)
    >>> y.shape
    (4, 65)
    """
    # Dimension key: H=height, W=width, D=spatial_dim, A=ambient_dim (D+1)

    x_space_HWD = x[..., 1:]  # (..., H, W, D) — drop time coordinate
    x_pooled_D = jnp.mean(x_space_HWD, axis=(-3, -2))  # (..., D)
    return spatial_to_hyperboloid(x_pooled_D, c_in=c, c_out=c, eps=eps)  # (..., D+1)

Example

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

conv = HypConv2DHyperboloid(
    manifold_module=Hyperboloid(), in_channels=16, out_channels=32,
    kernel_size=(3, 3), stride=(1, 1), rngs=nnx.Rngs(0),
)
x = jax.random.normal(jax.random.PRNGKey(1), (8, 28, 28, 16))
x_amb = jnp.concatenate([jnp.sqrt(jnp.sum(x**2, -1, keepdims=True) + 1.0), x], -1)  # (8,28,28,17)
output = conv(x_amb, c=1.0)
print(output.shape)  # (8, 28, 28, 32*9+1) — dimension grows via HCat