hyperbolix 0.4.0__tar.gz → 0.4.2__tar.gz

{hyperbolix-0.4.0 → hyperbolix-0.4.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: hyperbolix
-Version: 0.4.0
+Version: 0.4.2
 Summary: Hyperbolic Deep Learning in JAX
 Keywords: hyperbolic,deep-learning,jax,flax,machine-learning,riemannian,poincare,hyperboloid,neural-networks,geometric-deep-learning
 Author: Timo Klein, Thomas Lang
@@ -42,7 +42,7 @@ Pure JAX implementation of hyperbolic deep learning with manifold operations, ne
 ## Features
 
 - 🌐 **3 Manifolds**: Euclidean, Poincaré Ball, Hyperboloid
-- 🧠 **13+ Neural Network Layers**: Linear, convolutional (2D/3D), regression
+- 🧠 **13+ Neural Network Layers**: Linear, convolutional, regression
 - ⚡ **4 Hyperbolic Activations**: ReLU, Leaky ReLU, Tanh, Swish
 - 📈 **Riemannian Optimizers**: RAdam and RSGD with automatic manifold detection
 - 🚀 **Pure JAX/Flax NNX**: vmap-native API, JIT-compatible (10-100x speedup)

{hyperbolix-0.4.0 → hyperbolix-0.4.2}/README.md

@@ -12,7 +12,7 @@ Pure JAX implementation of hyperbolic deep learning with manifold operations, ne
 ## Features
 
 - 🌐 **3 Manifolds**: Euclidean, Poincaré Ball, Hyperboloid
-- 🧠 **13+ Neural Network Layers**: Linear, convolutional (2D/3D), regression
+- 🧠 **13+ Neural Network Layers**: Linear, convolutional, regression
 - ⚡ **4 Hyperbolic Activations**: ReLU, Leaky ReLU, Tanh, Swish
 - 📈 **Riemannian Optimizers**: RAdam and RSGD with automatic manifold detection
 - 🚀 **Pure JAX/Flax NNX**: vmap-native API, JIT-compatible (10-100x speedup)

{hyperbolix-0.4.0 → hyperbolix-0.4.2}/hyperbolix/nn_layers/__init__.py

@@ -19,7 +19,7 @@ from .hyperboloid_attention import (
     HyperbolicSoftmaxAttention,
     focus_transform,
 )
-from .hyperboloid_conv import FGGConv2D, HypConv2DHyperboloid, HypConv3DHyperboloid, LorentzConv2D
+from .hyperboloid_conv import FGGConv2D, HypConv2DHyperboloid, LorentzConv2D
 from .hyperboloid_core import build_spacelike_V, hrc, htc, lorentz_midpoint, lorentz_residual, spatial_to_hyperboloid
 from .hyperboloid_linear import FGGLinear, HTCLinear, HypLinearHyperboloidFHCNN
 from .hyperboloid_positional import HyperbolicRoPE, HypformerPositionalEncoding, hope
@@ -43,7 +43,6 @@ __all__ = [
     "HTCLinear",
     "HypConv2DHyperboloid",
     "HypConv2DPoincare",
-    "HypConv3DHyperboloid",
     "HypLinearHyperboloidFHCNN",
     "HypLinearPoincare",
     "HypLinearPoincarePP",

{hyperbolix-0.4.0 → hyperbolix-0.4.2}/hyperbolix/nn_layers/hyperboloid_conv.py

@@ -2,8 +2,8 @@
 
 Dimension key:
   B: batch size
-  H: output height        W: output width         Z: output depth (3D conv)
-  C: channels (in/out)    K: kernel elements (kh*kw or kd*kh*kw)
+  H: output height        W: output width
+  C: channels (in/out)    K: kernel elements (kh*kw)
   A: ambient dimension (in_channels or hcat output dim)
 """
 
@@ -18,7 +18,7 @@ from hyperbolix.manifolds.hyperboloid import Hyperboloid
 
 from ._helpers import validate_hyperboloid_manifold
 from .hyperboloid_core import hrc
-from .hyperboloid_linear import FGGLinear, HypLinearHyperboloidFHCNN
+from .hyperboloid_linear import _fgg_linear_forward, _fhcnn_forward, _get_effective_kernel
 
 
 class LorentzConv2D(nnx.Module):
@@ -230,17 +230,13 @@ class HypConv2DHyperboloid(nnx.Module):
         d = in_channels - 1  # Input manifold dimension
         hcat_out_ambient_dim = d * N + 1  # HCat output ambient dimension
 
-        # Create the linear transformation layer
-        # Input: hcat_out_ambient_dim, Output: out_channels
-        self.linear = HypLinearHyperboloidFHCNN(
-            manifold_module=self.manifold,
-            in_dim=hcat_out_ambient_dim,
-            out_dim=out_channels,
-            rngs=rngs,
-            input_space="manifold",  # HCat output is always on manifold
-            learnable_scale=False,
-            normalize=False,
+        # Trainable parameters — owned directly for flat parameter paths
+        bound = 0.02
+        self.kernel = nnx.Param(
+            jax.random.uniform(rngs.params(), (out_channels, hcat_out_ambient_dim), minval=-bound, maxval=bound)
         )
+        self.bias = nnx.Param(jnp.zeros((1, out_channels)))
+        self.scale = 2.3  # not learnable (matches FHCNN default)
 
     def _extract_patches(
         self,
@@ -322,7 +318,18 @@ class HypConv2DHyperboloid(nnx.Module):
         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 = self.linear(hcat_out_NA, c)  # (B*H*W, 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)
@@ -330,210 +337,6 @@ class HypConv2DHyperboloid(nnx.Module):
         return output_BHWC
 
 
-class HypConv3DHyperboloid(nnx.Module):
-    """
-    Hyperbolic 3D Convolutional Layer for Hyperboloid model.
-
-    This layer implements fully hyperbolic 3D convolution, extending the 2D approach
-    from "Fully Hyperbolic Convolutional Neural Networks for Computer Vision" to
-    volumetric data.
-
-    Computation steps:
-        1) Extract receptive field (kernel_size x 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
-    ----------
-    manifold_module : object
-        Class-based Hyperboloid manifold instance
-    in_channels : int
-        Number of input channels
-    out_channels : int
-        Number of output channels
-    kernel_size : int or tuple[int, int, int]
-        Size of the convolutional kernel (depth, height, width)
-    rngs : nnx.Rngs
-        Random number generators for parameter initialization
-    stride : int or tuple[int, int, int]
-        Stride of the convolution (default: 1)
-    padding : str
-        Padding mode, either 'SAME' or 'VALID' (default: '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.
-
-    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).
-    """
-
-    def __init__(
-        self,
-        manifold_module: Hyperboloid,
-        in_channels: int,
-        out_channels: int,
-        kernel_size: int | tuple[int, int, int],
-        *,
-        rngs: nnx.Rngs,
-        stride: int | tuple[int, int, int] = 1,
-        padding: str = "SAME",
-        input_space: str = "manifold",
-    ):
-        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
-
-        # Handle kernel_size as int or tuple
-        if isinstance(kernel_size, int):
-            self.kernel_size = (kernel_size, kernel_size, kernel_size)
-        else:
-            self.kernel_size = kernel_size
-
-        # Handle stride as int or tuple
-        if isinstance(stride, int):
-            self.stride = (stride, stride, stride)
-        else:
-            self.stride = stride
-
-        # Compute dimensions for the linear layer
-        # Receptive field: kernel_d x kernel_h x kernel_w voxels, each is a point in in_channels-dim ambient space
-        # HCat input: N = kernel_d x kernel_h x kernel_w points, each in in_channels ambient dimensions
-        # HCat output: ambient dim = (in_channels - 1) x N + 1
-        #            = (in_channels - 1) x (kernel_d x kernel_h x kernel_w) + 1
-        kernel_d, kernel_h, kernel_w = self.kernel_size
-        N = kernel_d * kernel_h * kernel_w  # Number of points in receptive field
-        d = in_channels - 1  # Input manifold dimension
-        hcat_out_ambient_dim = d * N + 1  # HCat output ambient dimension
-
-        # Create the linear transformation layer
-        # Input: hcat_out_ambient_dim, Output: out_channels
-        self.linear = HypLinearHyperboloidFHCNN(
-            manifold_module=self.manifold,
-            in_dim=hcat_out_ambient_dim,
-            out_dim=out_channels,
-            rngs=rngs,
-            input_space="manifold",  # HCat output is always on manifold
-            learnable_scale=False,
-            normalize=False,
-        )
-
-    def _extract_patches(
-        self,
-        x: Float[Array, "batch depth height width in_channels"],
-    ) -> Float[Array, "batch out_depth out_height out_width kernel_d kernel_h kernel_w in_channels"]:
-        """Extract patches (receptive fields) from the 3D input using optimized JAX primitives."""
-        batch, depth, height, width, in_channels = x.shape
-        kernel_d, kernel_h, kernel_w = self.kernel_size
-        stride_d, stride_h, stride_w = self.stride
-
-        # 1. Handle Padding (Manually, to ensure manifold validity)
-        if self.padding == "SAME":
-            out_depth = (depth + stride_d - 1) // stride_d
-            out_height = (height + stride_h - 1) // stride_h
-            out_width = (width + stride_w - 1) // stride_w
-
-            pad_d = max((out_depth - 1) * stride_d + kernel_d - depth, 0)
-            pad_h = max((out_height - 1) * stride_h + kernel_h - height, 0)
-            pad_w = max((out_width - 1) * stride_w + kernel_w - width, 0)
-
-            pad_front = pad_d // 2
-            pad_back = pad_d - pad_front
-            pad_top = pad_h // 2
-            pad_bottom = pad_h - pad_top
-            pad_left = pad_w // 2
-            pad_right = pad_w - pad_left
-
-            x = jnp.pad(
-                x,
-                ((0, 0), (pad_front, pad_back), (pad_top, pad_bottom), (pad_left, pad_right), (0, 0)),
-                mode="edge",
-            )
-
-        # 2. Extract Patches — output: (B, Z, H, W, C*kd*kh*kw)
-        patches_flat_BZHW_CKdKhKw = jax.lax.conv_general_dilated_patches(
-            lhs=x,
-            filter_shape=(kernel_d, kernel_h, kernel_w),
-            window_strides=(stride_d, stride_h, stride_w),
-            padding="VALID",
-            dimension_numbers=("NDHWC", "OIDHW", "NDHWC"),
-        )
-
-        # 3. Reshape to separate kernel dims and channels, then transpose
-        out_d, out_h, out_w = (
-            patches_flat_BZHW_CKdKhKw.shape[1],
-            patches_flat_BZHW_CKdKhKw.shape[2],
-            patches_flat_BZHW_CKdKhKw.shape[3],
-        )
-        # Flattened dim order from rhs_spec="OIDHW": (C, kd, kh, kw)
-        patches_BZHWCkdkhkw = patches_flat_BZHW_CKdKhKw.reshape(
-            batch, out_d, out_h, out_w, in_channels, kernel_d, kernel_h, kernel_w
-        )
-        patches_BZHWkdkhkwC = patches_BZHWCkdkhkw.transpose(0, 1, 2, 3, 5, 6, 7, 4)  # move C last
-
-        return patches_BZHWkdkhkwC
-
-    def __call__(
-        self,
-        x: Float[Array, "batch depth height width in_channels"],
-        c: float = 1.0,
-    ) -> Float[Array, "batch out_depth out_height out_width out_channels"]:
-        """
-        Forward pass through the hyperbolic 3D convolutional layer.
-
-        Parameters
-        ----------
-        x : Array of shape (batch, depth, height, width, in_channels)
-            Input feature map where each voxel is a point on the Hyperboloid manifold
-        c : float
-            Manifold curvature (default: 1.0)
-
-        Returns
-        -------
-        out : Array of shape (batch, out_depth, out_height, out_width, out_channels)
-            Output feature map on the Hyperboloid manifold
-        """
-        # Map to manifold if needed (static branch - JIT friendly)
-        if self.input_space == "tangent":
-            x_flat_NC = x.reshape(-1, x.shape[-1])  # (B*Z*H*W, C)
-            x_mapped_NC = jax.vmap(self.manifold.expmap_0, in_axes=(0, None))(x_flat_NC, c)
-            x = x_mapped_NC.reshape(x.shape)  # (B, Z, H, W, C)
-
-        # Extract patches: (B, Z, H, W, kd, kh, kw, C)
-        patches_BZHWkdkhkwC = self._extract_patches(x)
-        batch, out_d, out_h, out_w, kd, kh, kw, in_c = patches_BZHWkdkhkwC.shape
-
-        # Flatten batch+spatial for parallel processing: (B*Z*H*W, K, C)
-        patches_flat_NKC = patches_BZHWkdkhkwC.reshape(-1, kd * 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*Z*H*W, hcat_dim)
-
-        # Linear: (hcat_dim,) -> (out_channels,)
-        linear_out_NC = self.linear(hcat_out_NA, c)  # (B*Z*H*W, out_channels)
-
-        # Reshape back to spatial
-        output_BZHWC = linear_out_NC.reshape(batch, out_d, out_h, out_w, self.out_channels)
-
-        return output_BZHWC
-
-
 class FGGConv2D(nnx.Module):
     """Fast and Geometrically Grounded Lorentz 2D convolutional layer.
 
@@ -595,7 +398,7 @@ class FGGConv2D(nnx.Module):
         padding: str = "SAME",
         pad_mode: str = "origin",
         activation: Callable | None = None,
-        reset_params: str = "kaiming",
+        reset_params: str = "lorentz_kaiming",
         use_weight_norm: bool = False,
         init_bias: float = 0.5,
         eps: float = 1e-7,
@@ -627,17 +430,45 @@ class FGGConv2D(nnx.Module):
         kh, kw = self.kernel_size
         hcat_out_ambient = (in_channels - 1) * kh * kw + 1
 
-        # FGGLinear for channel transformation
-        self.fgg_linear = FGGLinear(
-            hcat_out_ambient,
-            out_channels,
-            rngs=rngs,
-            activation=activation,
-            reset_params=reset_params,
-            use_weight_norm=use_weight_norm,
-            init_bias=init_bias,
-            eps=eps,
-        )
+        # Trainable parameters — owned directly for flat parameter paths
+        if reset_params not in ("eye", "xavier", "kaiming", "lorentz_kaiming", "mlr"):
+            raise ValueError(
+                f"reset_params must be 'eye', 'xavier', 'kaiming', 'lorentz_kaiming', or 'mlr', got '{reset_params}'"
+            )
+
+        in_spatial = hcat_out_ambient - 1  # I
+        out_spatial = out_channels - 1  # O
+
+        self.activation = activation
+        self.use_weight_norm = use_weight_norm
+
+        # Initialize Euclidean weight U: (I, O)
+        # Reference computes std from ambient dimensions (hcat_out_ambient, out_channels)
+        key = rngs.params()
+        if reset_params == "eye":
+            U_init = 0.5 * jnp.eye(in_spatial, out_spatial)
+        elif reset_params == "xavier":
+            std = jnp.sqrt(1.0 / (hcat_out_ambient + out_channels))
+            U_init = jax.random.normal(key, (in_spatial, out_spatial)) * std
+        elif reset_params == "kaiming":
+            std = jnp.sqrt(2.0 / hcat_out_ambient)
+            U_init = jax.random.normal(key, (in_spatial, out_spatial)) * std
+        elif reset_params == "lorentz_kaiming":
+            std = jnp.sqrt(1.0 / hcat_out_ambient)
+            U_init = jax.random.normal(key, (in_spatial, out_spatial)) * std
+        else:  # mlr
+            std = jnp.sqrt(5.0 / hcat_out_ambient)
+            U_init = jax.random.normal(key, (in_spatial, out_spatial)) * std
+
+        # 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))  # (O,)
+        else:
+            self.kernel = nnx.Param(U_init)  # (I, O)
+
+        self.bias = nnx.Param(jnp.full((out_spatial,), init_bias))  # (O,)
 
     def _extract_patches(
         self,
@@ -717,8 +548,15 @@ class FGGConv2D(nnx.Module):
         # HCat: (K, C) -> (hcat_dim,) per patch
         hcat_out_NA = jax.vmap(self.manifold.hcat, in_axes=(0, None))(patches_flat_NKC, c)
 
-        # FGGLinear: (hcat_dim,) -> (out_channels,)
-        linear_out_NC = self.fgg_linear(hcat_out_NA, 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-0.4.0 → hyperbolix-0.4.2}/hyperbolix/nn_layers/hyperboloid_linear.py

@@ -25,6 +25,107 @@ from ._helpers import validate_hyperboloid_manifold
 from .hyperboloid_core import build_spacelike_V, htc
 
 
+def _fhcnn_forward(
+    x_BI: Float[Array, "batch in_dim"],
+    kernel_OI: Array,
+    bias_1O: Array,
+    manifold: Manifold,
+    c: float,
+    input_space: str,
+    activation: Callable[[Array], Array] | None,
+    normalize: bool,
+    scale_val: float | Array,
+    eps: float,
+) -> Float[Array, "batch out_dim"]:
+    """Pure-function FHCNN forward pass.
+
+    Used by both HypLinearHyperboloidFHCNN and HypConv2DHyperboloid.
+    """
+    # Map to manifold if needed (static branch - JIT friendly)
+    if input_space == "tangent":
+        x_BI = jax.vmap(manifold.expmap_0, in_axes=(0, None), out_axes=0)(x_BI, c)
+
+    # Apply activation if provided (static branch - JIT friendly)
+    if activation is not None:
+        x_BI = activation(x_BI)
+
+    # Linear transformation: (B, in_dim) -> (B, out_dim)
+    x_BO = jnp.einsum("bi,oi->bo", x_BI, kernel_OI) + bias_1O
+
+    # Split into time and space: x0 is first coord, x_rem is spatial
+    x0_B1 = x_BO[:, 0:1]  # (B, 1) -- time coordinate
+    x_rem_BD = x_BO[:, 1:]  # (B, D) where D = out_dim - 1
+
+    # Static branch - JIT friendly
+    if normalize:
+        x_rem_norm_B1 = jnp.linalg.norm(x_rem_BD, ord=2, axis=-1, keepdims=True)  # (B, 1)
+
+        # Learnable sigmoid scaling
+        scale_B1 = jnp.exp(scale_val) * jax.nn.sigmoid(x0_B1)  # (B, 1)
+
+        res0_B1 = jnp.sqrt(scale_B1**2 + 1 / c + eps)  # (B, 1)
+        res_rem_BD = scale_B1 * x_rem_BD / x_rem_norm_B1  # (B, D)
+
+        res_BA = jnp.concatenate([res0_B1, res_rem_BD], axis=-1)  # (B, A)
+
+        # Cast near-zero-norm vectors to origin
+        origin_time_B1 = jnp.sqrt(1 / c) * jnp.ones_like(res0_B1)
+        origin_BA = jnp.concatenate([origin_time_B1, jnp.zeros_like(res_rem_BD)], axis=-1)
+
+        mask_B1 = x_rem_norm_B1 <= 1e-5
+        res_BA = jnp.where(mask_B1, origin_BA, res_BA)
+    else:
+        # Reconstruct time from space: x0 = sqrt(||x_rem||^2 + 1/c)
+        res0_B1 = jnp.sqrt(jnp.sum(x_rem_BD**2, axis=-1, keepdims=True) + 1 / c)  # (B, 1)
+        res_BA = jnp.concatenate([res0_B1, x_rem_BD], axis=-1)  # (B, A)
+
+    return res_BA
+
+
+def _get_effective_kernel(
+    kernel: Array | None,
+    kernel_dir: Array | None,
+    kernel_scale: Array | None,
+    use_weight_norm: bool,
+    eps: float,
+) -> Array:
+    """Compute effective weight matrix, handling weight normalization."""
+    if use_weight_norm:
+        assert kernel_scale is not None and kernel_dir is not None
+        g_pos_O = jax.nn.softplus(kernel_scale)  # (O,) force positive magnitudes
+        v_norm_O = jnp.sqrt(jnp.sum(kernel_dir**2, axis=0) + eps)  # (O,)
+        return g_pos_O[None, :] * kernel_dir / v_norm_O[None, :]  # (I, O)
+    assert kernel is not None
+    return kernel
+
+
+def _fgg_linear_forward(
+    x_BAi: Float[Array, "batch in_features"],
+    U_IO: Array,
+    bias_O: Array,
+    c: float,
+    activation: Callable[[jax.Array], jax.Array] | None,
+    eps: float,
+) -> Float[Array, "batch out_features"]:
+    """Pure-function FGG forward: build_spacelike_V -> matmul -> activation -> time reconstruct."""
+    # Build V_mink from (kernel, bias) -- Minkowski metric absorbed
+    V_AiO = build_spacelike_V(U_IO, bias_O, c, eps)  # (Ai, O)
+    # Cast V to match input dtype (avoids float32/float64 scatter warnings)
+    V_AiO = V_AiO.astype(x_BAi.dtype)
+
+    # Minkowski inner products via matmul (metric in V)
+    z_BO = x_BAi @ V_AiO  # (B, O)
+
+    # Apply Euclidean activation (Lorentzian wrapping implicit via cancellation)
+    if activation is not None:
+        z_BO = activation(z_BO)
+
+    # Reconstruct hyperboloid point: spatial = z, time from constraint
+    y_0_B1 = jnp.sqrt(jnp.sum(z_BO**2, axis=-1, keepdims=True) + 1.0 / c)  # (B, 1)
+
+    return jnp.concatenate([y_0_B1, z_BO], axis=-1)  # (B, Ao)
+
+
 class HypLinearHyperboloidFHCNN(nnx.Module):
     """
     Fully Hyperbolic Convolutional Neural Networks fully connected layer (Hyperboloid model).
@@ -147,46 +248,19 @@ class HypLinearHyperboloidFHCNN(nnx.Module):
         res : Array of shape (batch, out_dim)
             Output on the Hyperboloid manifold
         """
-        # 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)
-
-        # Apply activation if provided (static branch - JIT friendly)
-        if self.activation is not None:
-            x = self.activation(x)
-
-        # Linear transformation: (B, in_dim) → (B, out_dim)
-        x_BO = jnp.einsum("bi,oi->bo", x, self.kernel) + self.bias
-
-        # Split into time and space: x0 is first coord, x_rem is spatial
-        x0_B1 = x_BO[:, 0:1]  # (B, 1) — time coordinate
-        x_rem_BD = x_BO[:, 1:]  # (B, D) where D = out_dim - 1
-
-        # Static branch - JIT friendly
-        if self.normalize:
-            x_rem_norm_B1 = jnp.linalg.norm(x_rem_BD, ord=2, axis=-1, keepdims=True)  # (B, 1)
-
-            # Learnable sigmoid scaling
-            scale_val = self.scale[...] if isinstance(self.scale, nnx.Param) else self.scale
-            scale_B1 = jnp.exp(scale_val) * jax.nn.sigmoid(x0_B1)  # (B, 1)
-
-            res0_B1 = jnp.sqrt(scale_B1**2 + 1 / c + self.eps)  # (B, 1)
-            res_rem_BD = scale_B1 * x_rem_BD / x_rem_norm_B1  # (B, D)
-
-            res_BA = jnp.concatenate([res0_B1, res_rem_BD], axis=-1)  # (B, A)
-
-            # Cast near-zero-norm vectors to origin
-            origin_time_B1 = jnp.sqrt(1 / c) * jnp.ones_like(res0_B1)
-            origin_BA = jnp.concatenate([origin_time_B1, jnp.zeros_like(res_rem_BD)], axis=-1)
-
-            mask_B1 = x_rem_norm_B1 <= 1e-5
-            res_BA = jnp.where(mask_B1, origin_BA, res_BA)
-        else:
-            # Reconstruct time from space: x₀ = sqrt(||x_rem||² + 1/c)
-            res0_B1 = jnp.sqrt(jnp.sum(x_rem_BD**2, axis=-1, keepdims=True) + 1 / c)  # (B, 1)
-            res_BA = jnp.concatenate([res0_B1, x_rem_BD], axis=-1)  # (B, A)
-
-        return res_BA
+        scale_val = self.scale[...] if isinstance(self.scale, nnx.Param) else self.scale
+        return _fhcnn_forward(
+            x,
+            self.kernel[...],
+            self.bias[...],
+            self.manifold,
+            c,
+            self.input_space,
+            self.activation,
+            self.normalize,
+            scale_val,
+            self.eps,
+        )
 
 
 class HTCLinear(nnx.Module):
@@ -392,27 +466,28 @@ class FGGLinear(nnx.Module):
         self.eps = eps
 
         # Initialize Euclidean weight U: (I, O)
+        # Reference computes std from ambient dimensions (in_features, out_features)
         key = rngs.params()
         if reset_params == "eye":
             U_init = 0.5 * jnp.eye(in_spatial, out_spatial)
         elif reset_params == "xavier":
-            std = jnp.sqrt(1.0 / (in_spatial + out_spatial))
+            std = jnp.sqrt(1.0 / (in_features + out_features))
             U_init = jax.random.normal(key, (in_spatial, out_spatial)) * std
         elif reset_params == "kaiming":
-            std = jnp.sqrt(2.0 / in_spatial)
+            std = jnp.sqrt(2.0 / in_features)
             U_init = jax.random.normal(key, (in_spatial, out_spatial)) * std
         elif reset_params == "lorentz_kaiming":
-            std = jnp.sqrt(1.0 / in_spatial)
+            std = jnp.sqrt(1.0 / in_features)
             U_init = jax.random.normal(key, (in_spatial, out_spatial)) * std
         else:  # mlr
-            std = jnp.sqrt(5.0 / in_spatial)
+            std = jnp.sqrt(5.0 / in_features)
             U_init = jax.random.normal(key, (in_spatial, out_spatial)) * std
 
         # Weight normalization: decompose kernel = softplus(kernel_scale) * kernel_dir / ||kernel_dir||
         if use_weight_norm:
             # Reference: kernel_dir from reset_params (normalized in forward), kernel_scale fixed magnitude
             self.kernel_dir = nnx.Param(U_init)  # (I, O) direction
-            g_init_val = jnp.sqrt(1.0 / (in_spatial + out_spatial))
+            g_init_val = jnp.sqrt(1.0 / (in_features + out_features))
             self.kernel_scale = nnx.Param(jnp.full((out_spatial,), g_init_val))  # (O,)
         else:
             self.kernel = nnx.Param(U_init)  # (I, O)
@@ -422,13 +497,13 @@ class FGGLinear(nnx.Module):
 
     def _get_kernel(self) -> jax.Array:
         """Return the effective weight matrix, handling weight normalization."""
-        if self.use_weight_norm:
-            v_IO = self.kernel_dir[...]  # (I, O)
-            g_O = self.kernel_scale[...]  # (O,)
-            g_pos_O = jax.nn.softplus(g_O)  # (O,) force positive magnitudes
-            v_norm_O = jnp.sqrt(jnp.sum(v_IO**2, axis=0) + self.eps)  # (O,)
-            return g_pos_O[None, :] * v_IO / v_norm_O[None, :]  # (I, O)
-        return self.kernel[...]  # (I, O)
+        return _get_effective_kernel(
+            self.kernel[...] if not self.use_weight_norm else None,
+            self.kernel_dir[...] if self.use_weight_norm else None,
+            self.kernel_scale[...] if self.use_weight_norm else None,
+            self.use_weight_norm,
+            self.eps,
+        )
 
     def __call__(
         self,
@@ -451,22 +526,5 @@ class FGGLinear(nnx.Module):
             Output points on the hyperboloid with curvature ``c``.
             Ao = out_features (ambient dimension).
         """
-        # 1. Get effective kernel (handle weight norm)
         U_IO = self._get_kernel()  # (I, O)
-
-        # 2. Build V_mink from (kernel, bias) — Minkowski metric absorbed
-        V_AiO = build_spacelike_V(U_IO, self.bias[...], c, self.eps)  # (Ai, O)
-        # Cast V to match input dtype (avoids float32/float64 scatter warnings)
-        V_AiO = V_AiO.astype(x_BAi.dtype)
-
-        # 3. Minkowski inner products via matmul (metric in V)
-        z_BO = x_BAi @ V_AiO  # (B, O)
-
-        # 4. Apply Euclidean activation (Lorentzian wrapping implicit via cancellation)
-        if self.activation is not None:
-            z_BO = self.activation(z_BO)
-
-        # 5. Reconstruct hyperboloid point: spatial = z, time from constraint
-        y_0_B1 = jnp.sqrt(jnp.sum(z_BO**2, axis=-1, keepdims=True) + 1.0 / c)  # (B, 1)
-
-        return jnp.concatenate([y_0_B1, z_BO], axis=-1)  # (B, Ao)
+        return _fgg_linear_forward(x_BAi, U_IO, self.bias[...], c, self.activation, self.eps)

{hyperbolix-0.4.0 → hyperbolix-0.4.2}/hyperbolix/nn_layers/hyperboloid_regression.py

@@ -174,13 +174,13 @@ class FGGLorentzMLR(nnx.Module):
         self.eps = eps
 
         # Hyperplane normals (spatial) and bias offsets
+        # Reference computes std from ambient dimension (in_features)
         key = rngs.params()
         if reset_params == "mlr":
-            # Reference: N(0, sqrt(5/I)) for spatial weights
-            std = jnp.sqrt(5.0 / in_spatial)
+            std = jnp.sqrt(5.0 / in_features)
             self.kernel = nnx.Param(jax.random.normal(key, (in_spatial, num_classes)) * std)
         else:  # default
-            stdv = 1.0 / jnp.sqrt(jnp.array(in_spatial, dtype=jnp.float32))
+            stdv = 1.0 / jnp.sqrt(jnp.array(in_features, dtype=jnp.float32))
             self.kernel = nnx.Param(jax.random.uniform(key, (in_spatial, num_classes), minval=-stdv, maxval=stdv))
         self.bias = nnx.Param(jnp.full((num_classes,), init_bias))
 

{hyperbolix-0.4.0 → hyperbolix-0.4.2}/hyperbolix/nn_layers/poincare_conv.py

@@ -25,7 +25,7 @@ from jaxtyping import Array, Float
 from hyperbolix.manifolds.poincare import Poincare
 
 from ._helpers import validate_poincare_manifold
-from .poincare_linear import HypLinearPoincarePP
+from .poincare_linear import _poincare_pp_forward
 
 
 class HypConv2DPoincare(nnx.Module):
@@ -150,25 +150,23 @@ class HypConv2DPoincare(nnx.Module):
         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 = beta_n / beta_ni
-
-        # HNN++ linear: in_dim = K^2 * C_in, out_dim = C_out
-        self.linear = HypLinearPoincarePP(
-            manifold_module=self.manifold,
-            in_dim=concat_dim,
-            out_dim=out_channels,
-            rngs=rngs,
-            input_space="manifold",  # input is on manifold after expmap_0
-            clamping_factor=clamping_factor,
-            smoothing_factor=smoothing_factor,
-        )
+        self.beta_scale = float(beta_n / beta_ni)
+
+        self.clamping_factor = clamping_factor
+        self.smoothing_factor = smoothing_factor
 
-        # Override weight init with identity initialization from reference
-        # (van Spengler et al. 2023): 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.
+        # 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:
-            eye = jnp.eye(out_channels, concat_dim)  # (C_out, K^2*C_in)
-            self.linear.kernel[...] = 0.5 * eye
+            # 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)
+        else:
+            std = 1.0 / jnp.sqrt(concat_dim)
+            kernel_init = jax.random.normal(rngs.params(), (out_channels, concat_dim)) * std
+        self.kernel = nnx.Param(kernel_init)
+        self.bias = nnx.Param(jnp.zeros((out_channels, 1)))
 
     def __call__(
         self,
@@ -228,7 +226,16 @@ class HypConv2DPoincare(nnx.Module):
         )  # (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 = self.linear(manifold_pts_NKC, c)
+        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)

{hyperbolix-0.4.0 → hyperbolix-0.4.2}/hyperbolix/nn_layers/poincare_linear.py

@@ -128,6 +128,40 @@ class HypLinearPoincare(nnx.Module):
         return res_BO
 
 
+def _poincare_pp_forward(
+    x_BI: Float[Array, "batch in_dim"],
+    kernel_OI: Array,
+    bias_O1: Array,
+    manifold: Poincare,
+    c: float,
+    input_space: str,
+    clamping_factor: float,
+    smoothing_factor: float,
+) -> Float[Array, "batch out_dim"]:
+    """Pure-function HNN++ forward pass.
+
+    Used by both HypLinearPoincarePP and HypConv2DPoincare.
+    """
+    # Map to manifold if needed (static branch - JIT friendly)
+    if input_space == "tangent":
+        x_BI = jax.vmap(manifold.expmap_0, in_axes=(0, None), out_axes=0)(x_BI, c)
+
+    # Compute multinomial linear regression
+    v = manifold.compute_mlr_pp(x_BI, kernel_OI, bias_O1, c, clamping_factor, smoothing_factor)
+
+    # Generalized linear transformation
+    sqrt_c = jnp.sqrt(c)
+    w_BO = sinh(sqrt_c * v) / sqrt_c
+    w2_B1 = jnp.sum(w_BO**2, axis=-1, keepdims=True)
+    denom_B1 = 1 + jnp.sqrt(1 + c * w2_B1)
+    res_BO = w_BO / denom_B1  # (B, 1) broadcasts over (B, O)
+
+    # Project results to the manifold
+    res_BO = jax.vmap(manifold.proj, in_axes=(0, None), out_axes=0)(res_BO, c)
+
+    return res_BO
+
+
 class HypLinearPoincarePP(nnx.Module):
     """
     Hyperbolic Neural Networks ++ fully connected layer (Poincaré ball model).
@@ -220,28 +254,13 @@ class HypLinearPoincarePP(nnx.Module):
         res : Array of shape (batch, out_dim)
             Output on the Poincaré ball manifold
         """
-        # 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
-        v = self.manifold.compute_mlr_pp(
+        return _poincare_pp_forward(
             x,
             self.kernel[...],
             self.bias[...],
+            self.manifold,
             c,
+            self.input_space,
             self.clamping_factor,
             self.smoothing_factor,
         )
-
-        # Generalized linear transformation
-        sqrt_c = jnp.sqrt(c)
-        w_BO = sinh(sqrt_c * v) / sqrt_c
-        w2_B1 = jnp.sum(w_BO**2, axis=-1, keepdims=True)
-        denom_B1 = 1 + jnp.sqrt(1 + c * w2_B1)
-        res_BO = w_BO / denom_B1  # (B, 1) broadcasts over (B, O)
-
-        # Project results to the manifold
-        res_BO = jax.vmap(self.manifold.proj, in_axes=(0, None), out_axes=0)(res_BO, c)
-
-        return res_BO

{hyperbolix-0.4.0 → hyperbolix-0.4.2}/pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "uv_build"
 
 [project]
 name = "hyperbolix"
-version = "0.4.0"
+version = "0.4.2"
 description = "Hyperbolic Deep Learning in JAX"
 authors = [
     { name="Timo Klein", email="timo.klein@univie.ac.at" },