Equimo 0.4.3__tar.gz → 0.4.5__tar.gz

{equimo-0.4.3/src/Equimo.egg-info → equimo-0.4.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: Equimo
-Version: 0.4.3
+Version: 0.4.5
 Summary: Implementation of popular vision models in Jax
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Developers

{equimo-0.4.3 → equimo-0.4.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "Equimo"
-version = "0.4.3"
+version = "0.4.5"
 description = "Implementation of popular vision models in Jax"
 readme = "README.md"
 requires-python = ">=3.11"

{equimo-0.4.3 → equimo-0.4.5/src/Equimo.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: Equimo
-Version: 0.4.3
+Version: 0.4.5
 Summary: Implementation of popular vision models in Jax
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Developers

equimo-0.4.5/src/equimo/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.4.5"

{equimo-0.4.3 → equimo-0.4.5}/src/equimo/layers/convolution.py

@@ -1,6 +1,6 @@
 import math
 import operator
-from typing import Callable, Literal, Optional, Sequence, Tuple
+from typing import Callable, Literal, Optional, Sequence, Tuple, Union
 
 import equinox as eqx
 import jax
@@ -11,6 +11,7 @@ from jaxtyping import Array, Float, PRNGKeyArray
 
 from equimo.layers.dropout import DropPathAdd
 from equimo.layers.norm import LayerScale
+from equimo.layers.squeeze_excite import SEModule
 from equimo.utils import make_divisible, nearest_power_of_2_divisor
 
 
@@ -1046,31 +1047,50 @@ class UIB(eqx.Module):
 
 
 class GenericGhostModule(eqx.Module):
-    """Modded module for the GhostNet v3 model.
+    """GhostNet v3-like module with GroupNorm and training-time branch fusion.
 
+    - Uses shared GroupNorm after linear summations (keeps GN's batch-size robustness).
+    - Preserves the linear pre-norm structure of the reference:
+      primary:   [skip (identity) + scale(1x1) + Σ conv_k] → GN → act
+      cheap:     [skip (identity, when ratio=2) + scale(1x1 depthwise) + Σ depthwise conv_k] → GN → act
+    - Provides test-time fusion into single Conv2d layers (still followed by GN+act).
 
-    It differs from the original implementation because it shares the same norm
-    ayers accross multiple parallel branches to allow using other norm layers
-    while still fusing convolutions at test-time.
+    Input/Output convention: (C, H, W).
     """
 
+    # Static configuration
     mode: Literal["original", "shortcut"] = eqx.field(static=True)
     out_channels: int = eqx.field(static=True)
     num_conv_branches: int = eqx.field(static=True)
+    kernel_size: int = eqx.field(static=True)
+    dw_size: int = eqx.field(static=True)
+    stride: int = eqx.field(static=True)
+    primary_has_skip: bool = eqx.field(static=True)
+    primary_has_scale: bool = eqx.field(static=True)
+    cheap_has_skip: bool = eqx.field(static=True)
+    cheap_has_scale: bool = eqx.field(static=True)
+
+    # Runtime flags
     inference: bool
 
+    # Inference
     primary_conv: eqx.nn.Conv2d
     cheap_operation: eqx.nn.Conv2d
 
+    # Training
     primary_rpr_conv: list[eqx.nn.Conv2d]
+    primary_rpr_scale: eqx.nn.Conv2d | eqx.nn.Identity
     primary_shared_norm: eqx.nn.GroupNorm
     primary_activation: Callable
-    cheap_rpr_scale: eqx.nn.Conv2d | eqx.nn.Identity
-    cheap_rpr_skip: eqx.nn.Conv2d | eqx.nn.Identity
+
     cheap_rpr_conv: list[eqx.nn.Conv2d]
+    cheap_rpr_scale: eqx.nn.Conv2d | eqx.nn.Identity
     cheap_shared_norm: eqx.nn.GroupNorm
     cheap_activation: Callable
+
     short_conv: eqx.nn.Identity | eqx.nn.Sequential
+    pool2: eqx.nn.AvgPool2d
+    gate_scale: float = eqx.field(static=True)
 
     def __init__(
         self,
@@ -1086,8 +1106,9 @@ class GenericGhostModule(eqx.Module):
         mode: Literal["original", "shortcut"] = "original",
         key: PRNGKeyArray,
     ):
-        key_primary, key_cheap, key_crs, key_skip, key_s1, key_s2, key_s3 = jr.split(
-            key, 7
+        # Key management
+        key_primary, key_cheap, key_pscale, key_cscale, key_s1, key_s2, key_s3 = (
+            jr.split(key, 7)
         )
         key_ps = jr.split(key_primary, num_conv_branches)
         key_cs = jr.split(key_cheap, num_conv_branches)
@@ -1095,10 +1116,22 @@ class GenericGhostModule(eqx.Module):
         init_channels = math.ceil(out_channels / ratio)
         new_channels = init_channels * (ratio - 1)
 
+        primary_has_skip = (in_channels == init_channels) and (stride == 1)
+        primary_has_scale = kernel_size > 1
+        cheap_has_skip = init_channels == new_channels
+        cheap_has_scale = dw_size > 1
+
         self.inference = False
         self.mode = mode
         self.num_conv_branches = num_conv_branches
         self.out_channels = out_channels
+        self.kernel_size = kernel_size  # MOD
+        self.dw_size = dw_size  # MOD
+        self.stride = stride  # MOD
+        self.primary_has_skip = primary_has_skip  # MOD
+        self.primary_has_scale = primary_has_scale  # MOD
+        self.cheap_has_skip = cheap_has_skip  # MOD
+        self.cheap_has_scale = cheap_has_scale  # MOD
 
         # Those are actually placeholders, updated at each epoch, only used at inference time
         self.primary_conv = eqx.nn.Conv2d(
@@ -1107,20 +1140,22 @@ class GenericGhostModule(eqx.Module):
             kernel_size=kernel_size,
             stride=stride,
             padding=kernel_size // 2,
+            use_bias=False,
             key=key_primary,
         )
         self.cheap_operation = eqx.nn.Conv2d(
             in_channels=init_channels,
             out_channels=new_channels,
             kernel_size=dw_size,
-            stride=stride,
+            stride=1,  # MOD
             padding=dw_size // 2,
-            groups=1,
+            groups=init_channels,
+            use_bias=False,
             key=key_cheap,
         )
 
+        # Primary training branches
         init_num_groups = nearest_power_of_2_divisor(init_channels, 32)
-        # TODO: test with some dropout?
         self.primary_rpr_conv = [
             eqx.nn.Conv2d(
                 in_channels=in_channels,
@@ -1128,44 +1163,56 @@ class GenericGhostModule(eqx.Module):
                 kernel_size=kernel_size,
                 stride=stride,
                 padding=kernel_size // 2,
+                use_bias=False,
                 key=key_ps[i],
             )
             for i in range(num_conv_branches)
         ]
+        self.primary_rpr_scale = (
+            eqx.nn.Conv2d(
+                in_channels=in_channels,
+                out_channels=init_channels,
+                kernel_size=1,
+                stride=1,
+                padding=0,
+                use_bias=False,
+                key=key_pscale,
+            )
+            if primary_has_scale
+            else eqx.nn.Identity()
+        )
         self.primary_shared_norm = eqx.nn.GroupNorm(init_num_groups, init_channels)
         self.primary_activation = act_layer
 
-        self.cheap_rpr_scale = eqx.nn.Conv2d(
-            in_channels=init_channels,
-            out_channels=new_channels,
-            kernel_size=1,
-            stride=1,
-            padding=0,
-            groups=1,
-            key=key_crs,
-        )
-        self.cheap_rpr_skip = eqx.nn.Conv2d(
-            init_channels,
-            new_channels,
-            kernel_size=1,
-            stride=stride,
-            padding=0,
-            groups=1,
-            key=key_skip,
-        )
+        # Cheap training branches (depthwise)
+        newchannels_num_groups = nearest_power_of_2_divisor(new_channels, 32)
         self.cheap_rpr_conv = [
             eqx.nn.Conv2d(
                 in_channels=init_channels,
                 out_channels=new_channels,
                 kernel_size=dw_size,
-                stride=stride,
+                stride=1,
                 padding=dw_size // 2,
-                groups=1,
+                groups=init_channels,
+                use_bias=False,
                 key=key_cs[i],
             )
             for i in range(self.num_conv_branches)
         ]
-        newchannels_num_groups = nearest_power_of_2_divisor(new_channels, 32)
+        self.cheap_rpr_scale = (
+            eqx.nn.Conv2d(
+                in_channels=init_channels,
+                out_channels=new_channels,
+                kernel_size=1,
+                stride=1,
+                padding=0,
+                groups=init_channels,
+                use_bias=False,
+                key=key_cscale,
+            )
+            if cheap_has_scale
+            else eqx.nn.Identity()
+        )
         self.cheap_shared_norm = eqx.nn.GroupNorm(newchannels_num_groups, new_channels)
         self.cheap_activation = act_layer
 
@@ -1211,29 +1258,39 @@ class GenericGhostModule(eqx.Module):
             else eqx.nn.Identity()
         )
 
-    def training_features(self, x):
-        x1 = jax.tree_util.tree_reduce(
-            operator.add, [conv(x) for conv in self.primary_rpr_conv]
+        self.pool2 = eqx.nn.AvgPool2d(
+            kernel_size=2, stride=2, padding=0, use_ceil=False
         )
-        x1 = self.primary_activation(self.primary_shared_norm(x1))
+        self.gate_scale = 1.0
 
-        cheap_branches = [self.cheap_rpr_skip(x1), self.cheap_rpr_scale(x1)] + [
-            conv(x1) for conv in self.cheap_rpr_conv
-        ]
-
-        x2 = jax.tree_util.tree_reduce(operator.add, cheap_branches)
-        x2 = self.cheap_activation(self.cheap_shared_norm(x2))
+    def training_features(self, x):
+        # Primary path pre-norm linear sum
+        terms = []
+        if self.primary_has_skip:
+            terms.append(x)  # identity
+        if not isinstance(self.primary_rpr_scale, eqx.nn.Identity):
+            terms.append(self.primary_rpr_scale(x))
+        terms.extend(conv(x) for conv in self.primary_rpr_conv)
+        x1_sum = jax.tree_util.tree_reduce(operator.add, terms)
+        x1 = self.primary_activation(self.primary_shared_norm(x1_sum))
+
+        # Cheap path pre-norm linear sum
+        cheap_terms = []
+        if self.cheap_has_skip:
+            cheap_terms.append(x1)  # identity
+        if not isinstance(self.cheap_rpr_scale, eqx.nn.Identity):
+            cheap_terms.append(self.cheap_rpr_scale(x1))
+        cheap_terms.extend(conv(x1) for conv in self.cheap_rpr_conv)
+        x2_sum = jax.tree_util.tree_reduce(operator.add, cheap_terms)
+        x2 = self.cheap_activation(self.cheap_shared_norm(x2_sum))
 
         out = jnp.concatenate([x1, x2], axis=0)
-
         return out
 
     def inference_features(self, x):
         x1 = self.primary_activation(self.primary_shared_norm(self.primary_conv(x)))
         x2 = self.cheap_activation(self.cheap_shared_norm(self.cheap_operation(x1)))
-
         out = jnp.concatenate([x1, x2], axis=0)
-
         return out
 
     def __call__(
@@ -1242,24 +1299,14 @@ class GenericGhostModule(eqx.Module):
         key: PRNGKeyArray,
         inference: Optional[bool] = None,
     ):
-        # key_p, key_c = jr.split(key, 2)
-        # key_ps = jr.split(key_p, self.num_conv_branches)
-        # key_cs = jr.split(key_c, self.num_conv_branches)
-
-        if self.inference:
-            out = self.inference_features(x)
-        else:
-            out = self.training_features(x)
+        use_inference = self.inference if inference is None else inference
+        out = self.inference_features(x) if use_inference else self.training_features(x)
 
         if self.mode == "shortcut":
-            res = self.short_conv(reduce(x, "c (h 2) (w 2) -> c h w", reduction="mean"))
+            res = self.short_conv(self.pool2(x))
             gating_signal = jax.image.resize(
-                image=jax.nn.sigmoid(res),
-                shape=(
-                    res.shape[0],
-                    out.shape[1],
-                    out.shape[2],
-                ),  # (channels, height, width)
+                image=jax.nn.sigmoid(res / self.gate_scale),
+                shape=(res.shape[0], out.shape[1], out.shape[2]),
                 method="nearest",
             )
             out = out[: self.out_channels, :, :] * gating_signal
@@ -1267,139 +1314,681 @@ class GenericGhostModule(eqx.Module):
         return out
 
 
-def update_fused_ghostmodule(
-    module: GenericGhostModule,
-) -> GenericGhostModule:
-    """
-    Fuses the weights of parallel training branches into single inference convolutions.
+class GhostBottleneck(eqx.Module):
+    """Ghost bottleneck with optional SE and re-parameterizable depthwise stage."""
 
-    This function should be called after training and before switching the model
-    to inference mode. It operates on a single GenericGhostModule.
+    # Static config
+    stride: int = eqx.field(static=True)
+    dw_kernel_size: int = eqx.field(static=True)
+    use_shortcut_mode_in_ghost1: bool = eqx.field(static=True)
+    allow_identity_residual: bool = eqx.field(static=True)
 
-    Args:
-        module: An instance of GenericGhostModule with trained weights.
+    inference: bool
 
-    Returns:
-        A new GenericGhostModule with the `primary_conv` and `cheap_operation`
-        weights updated for inference.
-    """
-    # Aggregate weights
-    primary_weights = jax.tree_util.tree_reduce(
-        operator.add, [conv.weight for conv in module.primary_rpr_conv]
-    )
+    ghost1: "GenericGhostModule"
+    ghost2: "GenericGhostModule"
+
+    dw_conv: eqx.nn.Conv2d | eqx.nn.Identity
+    dw_rpr_conv: list[eqx.nn.Conv2d]  # depthwise conv branches (no bias)
+    dw_rpr_scale: eqx.nn.Conv2d | eqx.nn.Identity  # optional 1x1 depthwise (no bias)
+    dw_shared_norm: eqx.nn.GroupNorm | eqx.nn.Identity
+
+    se: eqx.Module | eqx.nn.Identity
+
+    shortcut: eqx.nn.Sequential | eqx.nn.Identity
 
-    # Aggregate biases (if they exist)
-    primary_biases = None
-    if module.primary_rpr_conv[0].bias is not None:
-        primary_biases = jax.tree_util.tree_reduce(
-            operator.add, [conv.bias for conv in module.primary_rpr_conv]
+    def __init__(
+        self,
+        in_channels: int,
+        mid_channels: int,
+        out_channels: int,
+        *,
+        dw_kernel_size: int = 3,
+        stride: int = 1,
+        act_layer: Callable = jax.nn.relu,
+        se_ratio: float = 0.0,
+        layer_id: Optional[int] = None,
+        use_shortcut_mode_in_ghost1: bool = True,
+        allow_identity_residual: bool = True,
+        key: PRNGKeyArray,
+    ):
+        self.stride = stride
+        self.dw_kernel_size = dw_kernel_size
+        self.inference = False
+        self.allow_identity_residual = allow_identity_residual
+
+        self.use_shortcut_mode_in_ghost1 = use_shortcut_mode_in_ghost1
+
+        k_g1, k_g2, k_dw_main, k_dw_scale, k_sc1, k_sc2 = jr.split(key, 6)
+        k_dw_list = jr.split(k_dw_main, 3)
+
+        # ghost1 (expansion)
+        self.ghost1 = GenericGhostModule(
+            in_channels=in_channels,
+            out_channels=mid_channels,
+            kernel_size=1,
+            ratio=2,
+            dw_size=dw_kernel_size,
+            stride=1 if self.use_shortcut_mode_in_ghost1 else 1,
+            act_layer=act_layer,
+            num_conv_branches=3,
+            mode="shortcut" if self.use_shortcut_mode_in_ghost1 else "original",
+            key=k_g1,
         )
 
-    # Create the new fused primary convolution layer
-    fused_primary_conv = eqx.tree_at(
-        lambda conv: (conv.weight, conv.bias),
-        module.primary_conv,
-        (primary_weights, primary_biases),
-    )
+        # Depthwise stage (only if stride > 1)
+        if stride > 1:
+            # Training-time branches (depthwise, no bias); no activation; shared GN after sum
+            self.dw_rpr_conv = [
+                eqx.nn.Conv2d(
+                    in_channels=mid_channels,
+                    out_channels=mid_channels,
+                    kernel_size=dw_kernel_size,
+                    stride=stride,
+                    padding=(dw_kernel_size - 1) // 2,
+                    groups=mid_channels,
+                    use_bias=False,
+                    key=k_dw_list[i],
+                )
+                for i in range(3)
+            ]
+            # Optional scale branch (1x1, depthwise, stride=stride)
+            self.dw_rpr_scale = (
+                eqx.nn.Conv2d(
+                    in_channels=mid_channels,
+                    out_channels=mid_channels,
+                    kernel_size=1,
+                    stride=stride,
+                    padding=0,
+                    groups=mid_channels,
+                    use_bias=False,
+                    key=k_dw_scale,
+                )
+                if dw_kernel_size > 1
+                else eqx.nn.Identity()
+            )
+            # Shared GroupNorm
+            num_groups_dw = nearest_power_of_2_divisor(mid_channels, 32)
+            self.dw_shared_norm = eqx.nn.GroupNorm(num_groups_dw, mid_channels)
 
-    # The cheap operation fuses the skip, scale, and depthwise-like conv branches.
-    target_kernel_size = module.cheap_rpr_conv[0].weight.shape[-2:]
+            # Inference fused depthwise conv (weights filled by update)
+            self.dw_conv = eqx.nn.Conv2d(
+                in_channels=mid_channels,
+                out_channels=mid_channels,
+                kernel_size=dw_kernel_size,
+                stride=stride,
+                padding=(dw_kernel_size - 1) // 2,
+                groups=mid_channels,
+                use_bias=False,
+                key=k_dw_main,
+            )
+        else:
+            # No depthwise stage when stride == 1
+            self.dw_rpr_conv = []
+            self.dw_rpr_scale = eqx.nn.Identity()
+            self.dw_shared_norm = eqx.nn.Identity()
+            self.dw_conv = eqx.nn.Identity()
+
+        # SE
+        if se_ratio is not None and se_ratio > 0.0:
+            # Use provided SEModule with rd_ratio=se_ratio
+            self.se = SEModule(dim=mid_channels, rd_ratio=se_ratio, key=k_sc1)
+        else:
+            self.se = eqx.nn.Identity()
 
-    def pad_kernel_to_target(kernel: Array, target_size: tuple[int, int]) -> Array:
-        current_size = kernel.shape[-2:]
-        if current_size == target_size:
-            return kernel
+        # ghost2 (projection)
+        self.ghost2 = GenericGhostModule(
+            in_channels=mid_channels,
+            out_channels=out_channels,
+            kernel_size=1,
+            ratio=2,
+            dw_size=dw_kernel_size,
+            stride=1,
+            act_layer=act_layer,
+            num_conv_branches=3,
+            mode="original",
+            key=k_g2,
+        )
+
+        # Shortcut
+        if (in_channels == out_channels) and (stride == 1):
+            self.shortcut = eqx.nn.Identity()
+        else:
+            num_groups_sc_in = nearest_power_of_2_divisor(in_channels, 32)
+            num_groups_sc_out = nearest_power_of_2_divisor(out_channels, 32)
+            self.shortcut = eqx.nn.Sequential(
+                [
+                    eqx.nn.Conv2d(
+                        in_channels=in_channels,
+                        out_channels=in_channels,
+                        kernel_size=dw_kernel_size,
+                        stride=stride,
+                        padding=(dw_kernel_size - 1) // 2,
+                        groups=in_channels,
+                        use_bias=False,
+                        key=k_sc1,
+                    ),
+                    eqx.nn.GroupNorm(num_groups_sc_in, in_channels),
+                    eqx.nn.Conv2d(
+                        in_channels=in_channels,
+                        out_channels=out_channels,
+                        kernel_size=1,
+                        stride=1,
+                        padding=0,
+                        use_bias=False,
+                        key=k_sc2,
+                    ),
+                    eqx.nn.GroupNorm(num_groups_sc_out, out_channels),
+                ]
+            )
 
-        pad_h = (target_size[0] - current_size[0]) // 2
-        pad_w = (target_size[1] - current_size[1]) // 2
+    def __call__(
+        self,
+        x: Float[Array, "channels height width"],
+        *,
+        key: PRNGKeyArray,
+        inference: Optional[bool] = None,
+    ) -> Float[Array, "channels height width"]:
+        use_inference = self.inference if inference is None else inference
+        k_g1, k_g2 = jr.split(key, 2)
 
-        # JAX padding format: ((before, after), (before, after), ...)
-        padding_config = [
-            (0, 0),  # out_channels
-            (0, 0),  # in_channels
-            (pad_h, pad_h),  # height
-            (pad_w, pad_w),  # width
-        ]
-        return jnp.pad(kernel, padding_config, "constant", constant_values=0)
+        residual = x
 
-    # Pad the 1x1 kernels from the scale and skip branches.
-    padded_scale_weight = pad_kernel_to_target(
-        module.cheap_rpr_scale.weight, target_kernel_size
-    )
-    padded_skip_weight = pad_kernel_to_target(
-        module.cheap_rpr_skip.weight, target_kernel_size
+        x = self.ghost1(x, key=k_g1, inference=use_inference)
+
+        if self.stride > 1:
+            if use_inference:
+                x = self.dw_shared_norm(self.dw_conv(x))
+            else:
+                terms = []
+                if not isinstance(self.dw_rpr_scale, eqx.nn.Identity):
+                    terms.append(self.dw_rpr_scale(x))
+                terms.extend(conv(x) for conv in self.dw_rpr_conv)
+                x = self.dw_shared_norm(jax.tree_util.tree_reduce(operator.add, terms))
+
+        x = self.se(x)
+
+        x = self.ghost2(x, key=k_g2, inference=use_inference)
+
+        if not isinstance(self.shortcut, eqx.nn.Identity):
+            x += self.shortcut(residual)
+        elif self.allow_identity_residual:
+            x += residual
+
+        return x
+
+
+def _pad_kernel_to_target(kernel: Array, target_size: tuple[int, int]) -> Array:
+    kh, kw = kernel.shape[-2], kernel.shape[-1]
+    th, tw = target_size
+    if (kh, kw) == (th, tw):
+        return kernel
+    pad_h = (th - kh) // 2
+    pad_w = (tw - kw) // 2
+    padding_config = [
+        (0, 0),  # out_channels
+        (0, 0),  # in_channels_per_group
+        (pad_h, pad_h),  # height
+        (pad_w, pad_w),  # width
+    ]
+    return jnp.pad(kernel, padding_config, mode="constant", constant_values=0)
+
+
+def _make_identity_kernel_standard(out_ch: int, in_ch: int, kh: int, kw: int) -> Array:
+    """Identity kernel for standard conv (groups=1). Only valid if out_ch == in_ch."""
+    assert out_ch == in_ch, "Standard identity requires out_ch == in_ch."
+    k = jnp.zeros((out_ch, in_ch, kh, kw))
+    cy, cx = kh // 2, kw // 2
+    idx = jnp.arange(out_ch)
+    k = k.at[idx, idx, cy, cx].set(1.0)
+    return k
+
+
+def _make_identity_kernel_depthwise(out_ch: int, in_ch: int, kh: int, kw: int) -> Array:
+    """Identity kernel for depthwise conv with channel multiplier m = out_ch // in_ch."""
+    assert out_ch % in_ch == 0, "Depthwise identity requires out_ch multiple of in_ch."
+    m = out_ch // in_ch  # channel multiplier
+    k = jnp.zeros((out_ch, 1, kh, kw))
+    cy, cx = kh // 2, kw // 2
+    # For identity skip, we only use this when m==1 (ratio=2). Keep generic for clarity.
+    for c in range(in_ch):
+        for r in range(m):
+            k = k.at[c * m + r, 0, cy, cx].set(1.0)
+    return k
+
+
+def _update_ghostmodule(module: GenericGhostModule) -> GenericGhostModule:
+    """
+    Fuses the training branches into the inference convolutions (weights only).
+    Keeps GroupNorms; no bias is fused (bias is redundant with GN's affine).
+    """
+    if not isinstance(module, GenericGhostModule):
+        return module
+
+    # Fuse primary
+    primary_w = jax.tree_util.tree_reduce(
+        operator.add, [conv.weight for conv in module.primary_rpr_conv]
     )
+    if not isinstance(module.primary_rpr_scale, eqx.nn.Identity):
+        target_k = module.primary_conv.weight.shape[-2:]
+        scale_w = _pad_kernel_to_target(module.primary_rpr_scale.weight, target_k)
+        primary_w = primary_w + scale_w
+    if module.primary_has_skip:
+        kh, kw = module.primary_conv.weight.shape[-2:]
+        id_w = _make_identity_kernel_standard(
+            module.primary_conv.out_channels, module.primary_conv.in_channels, kh, kw
+        )
+        primary_w = primary_w + id_w
+    fused_primary = eqx.tree_at(lambda c: c.weight, module.primary_conv, primary_w)
 
-    # Sum the main cheap convolution weights.
-    cheap_conv_weights = jax.tree_util.tree_reduce(
+    # Fuse cheap (depthwise)
+    cheap_w = jax.tree_util.tree_reduce(
         operator.add, [conv.weight for conv in module.cheap_rpr_conv]
     )
+    if not isinstance(module.cheap_rpr_scale, eqx.nn.Identity):
+        target_k = module.cheap_operation.weight.shape[-2:]
+        scale_w = _pad_kernel_to_target(module.cheap_rpr_scale.weight, target_k)
+        cheap_w = cheap_w + scale_w
+    if module.cheap_has_skip:
+        init_ch = module.primary_conv.out_channels
+        new_ch = module.cheap_operation.out_channels
+        kh, kw = module.cheap_operation.weight.shape[-2:]
+        if new_ch == init_ch:
+            id_w = _make_identity_kernel_depthwise(new_ch, init_ch, kh, kw)
+            cheap_w = cheap_w + id_w
+    fused_cheap = eqx.tree_at(lambda c: c.weight, module.cheap_operation, cheap_w)
+
+    # Install fused convs
+    module = eqx.tree_at(lambda m: m.primary_conv, module, fused_primary)
+    module = eqx.tree_at(lambda m: m.cheap_operation, module, fused_cheap)
+
+    return module
 
-    # Sum all weights: main convs + padded scale + padded skip.
-    fused_cheap_weights = cheap_conv_weights + padded_scale_weight + padded_skip_weight
 
-    fused_cheap_biases = None
-    if module.cheap_rpr_conv[0].bias is not None:
-        cheap_conv_biases = jax.tree_util.tree_reduce(
-            operator.add, [conv.bias for conv in module.cheap_rpr_conv]
+def _finalize_ghostmodule(module: "GenericGhostModule") -> "GenericGhostModule":
+    """Finalize a GenericGhostModule for inference."""
+    if not isinstance(module, GenericGhostModule):
+        return module
+    fused = _update_ghostmodule(module)
+    # Drop training-only branches and switch to inference
+    fused = eqx.tree_at(lambda m: m.primary_rpr_conv, fused, list())
+    fused = eqx.tree_at(lambda m: m.cheap_rpr_conv, fused, list())
+    fused = eqx.tree_at(lambda m: m.primary_rpr_scale, fused, eqx.nn.Identity())
+    fused = eqx.tree_at(lambda m: m.cheap_rpr_scale, fused, eqx.nn.Identity())
+    fused = eqx.tree_at(lambda m: m.inference, fused, True)
+    return fused
+
+
+def _update_ghostbottleneck(module: "GhostBottleneck") -> "GhostBottleneck":
+    """Fuse depthwise stage and nested GhostModules for a GhostBottleneck."""
+    if not isinstance(module, GhostBottleneck):
+        return module
+
+    # Fuse depthwise stage (if present)
+    if module.stride > 1 and len(module.dw_rpr_conv) > 0:
+        dw_w = jax.tree_util.tree_reduce(
+            operator.add, [conv.weight for conv in module.dw_rpr_conv]
         )
-        scale_bias = module.cheap_rpr_scale.bias
-        skip_bias = module.cheap_rpr_skip.bias
-        fused_cheap_biases = cheap_conv_biases + scale_bias + skip_bias
-
-    fused_cheap_operation = eqx.tree_at(
-        lambda conv: (conv.weight, conv.bias),
-        module.cheap_operation,
-        (fused_cheap_weights, fused_cheap_biases),
-    )
+        if not isinstance(module.dw_rpr_scale, eqx.nn.Identity):
+            target_k = module.dw_conv.weight.shape[-2:]
+            scale_w = _pad_kernel_to_target(module.dw_rpr_scale.weight, target_k)
+            dw_w = dw_w + scale_w
+        fused_dw = eqx.tree_at(lambda c: c.weight, module.dw_conv, dw_w)
+        module = eqx.tree_at(lambda m: m.dw_conv, module, fused_dw)
+
+    # Update nested GhostModules
+    module = eqx.tree_at(lambda m: m.ghost1, module, _update_ghostmodule(module.ghost1))
+    module = eqx.tree_at(lambda m: m.ghost2, module, _update_ghostmodule(module.ghost2))
+    return module
 
-    module = eqx.tree_at(lambda m: m.primary_conv, module, fused_primary_conv)
-    module = eqx.tree_at(lambda m: m.cheap_operation, module, fused_cheap_operation)
 
-    return module
+def _finalize_ghostbottleneck(module: "GhostBottleneck") -> "GhostBottleneck":
+    """Finalize a GhostBottleneck for inference."""
+    if not isinstance(module, GhostBottleneck):
+        return module
+
+    fused = _update_ghostbottleneck(module)
+
+    # Drop depthwise training branches
+    if fused.stride > 1:
+        fused = eqx.tree_at(lambda m: m.dw_rpr_conv, fused, list())
+        fused = eqx.tree_at(lambda m: m.dw_rpr_scale, fused, eqx.nn.Identity())
+
+    # Finalize nested GhostModules
+    fused = eqx.tree_at(lambda m: m.ghost1, fused, _finalize_ghostmodule(fused.ghost1))
+    fused = eqx.tree_at(lambda m: m.ghost2, fused, _finalize_ghostmodule(fused.ghost2))
+
+    # Switch to inference
+    fused = eqx.tree_at(lambda m: m.inference, fused, True)
+    return fused
 
 
-def update_ghostnet(model: eqx.Module) -> eqx.Module:
+def update_ghostnet(
+    model: eqx.Module | Union["GenericGhostModule", "GhostBottleneck"],
+) -> eqx.Module | Union["GenericGhostModule", "GhostBottleneck"]:
     """
-    Update the inference layers of a GhostNet-like model.
-    Useful for intermediary evals.
+    Recursively fuse training branches for both GenericGhostModule and GhostBottleneck.
+    Keeps GroupNorm layers; no bias fusion.
     """
 
-    def _update_leaf(module: GenericGhostModule) -> GenericGhostModule:
-        if not isinstance(module, GenericGhostModule):
-            return module
-        return update_fused_ghostmodule(module)
+    def _update_leaf(m):
+        if isinstance(m, GenericGhostModule):
+            return _update_ghostmodule(m)
+        if isinstance(m, GhostBottleneck):
+            return _update_ghostbottleneck(m)
+        return m
 
-    is_g_module = lambda m: isinstance(m, GenericGhostModule)
-    return jax.tree_util.tree_map(_update_leaf, model, is_leaf=is_g_module)
+    is_leaf = lambda m: isinstance(m, (GenericGhostModule, GhostBottleneck))
+    return jax.tree_util.tree_map(_update_leaf, model, is_leaf=is_leaf)
 
 
-def finalize_ghostnet(model: eqx.Module) -> eqx.Module:
+def finalize_ghostnet(
+    model: eqx.Module | Union["GenericGhostModule", "GhostBottleneck"],
+) -> eqx.Module | Union["GenericGhostModule", "GhostBottleneck"]:
     """
-    Finalizes a trained GhostNet-like model for inference.
+    Recursively finalize both GenericGhostModule and GhostBottleneck for inference:
+    - Fuse training branches.
+    - Remove training-only branches.
+    - Switch to inference path.
+    """
+
+    def _finalize_leaf(m):
+        if isinstance(m, GenericGhostModule):
+            return _finalize_ghostmodule(m)
+        if isinstance(m, GhostBottleneck):
+            return _finalize_ghostbottleneck(m)
+        return m
+
+    is_leaf = lambda m: isinstance(m, (GenericGhostModule, GhostBottleneck))
+    return jax.tree_util.tree_map(_finalize_leaf, model, is_leaf=is_leaf)
+
+
+class PartialConv2d(eqx.Module):
+    """Partial 2D convolution on the channel dimension.
+
+    This layer applies a standard 2D convolution only to the first `C // n_dim`
+    input channels and leaves the remaining channels untouched (identity). It
+    follows the "partial convolution" idea used to increase throughput by
+    reducing compute on a subset of channels while preserving overall tensor
+    shape, as explored in [1].
+
+    References:
+      [1]. Chen et al., "Run, Don't Walk: Chasing Higher FLOPS for Faster Neural
+           Networks" [arXiv:2303.03667](https://arxiv.org/abs/2303.03667).
+
+    Implementation details:
+    - Let `C` be `in_channels` and `c = C // n_dim`. Only the first `c`
+      channels are convolved with a `Conv2d(c, c, ...)`. The remaining
+      `C - c` channels are forwarded unchanged.
+    - The forward pass uses a functional "update-slice" pattern
+      (`x.at[:c, ...].set(y1)`), which compiles to an efficient
+      `dynamic_update_slice` under XLA. With JIT buffer donation, this can be
+      performed in-place by the compiler.
+    - Spatial dimensions must be preserved by the convolution so that the
+      updated slice matches the input slice shape (e.g., use `stride=1` and
+      `padding="SAME"`). If spatial dimensions change, the slice update will
+      fail with a shape error.
+
+    Attributes
+    - dim: Number of channels to be convolved, computed as `in_channels // n_dim`.
+           This is treated as a static field for compilation stability.
+    - conv: The underlying `eqx.nn.Conv2d(c, c, ...)` applied to the first `c`
+            channels.
+
+    Notes
+    - FLOPs reduction is approximately `c / C = 1 / n_dim` relative to a full
+      convolution with the same kernel.
+    """
+
+    dim: int = eqx.field(static=True)
+
+    conv: eqx.nn.Conv2d
+
+    def __init__(
+        self,
+        in_channels: int,
+        n_dim: int,
+        *,
+        key: PRNGKeyArray,
+        kernel_size: int = 3,
+        padding: str | int = "SAME",
+        use_bias: bool = False,
+        **kwargs,
+    ):
+        """
+        Initialize a PartialConv2d layer.
+
+        Parameters
+        - in_channels: Total number of input channels `C`.
+        - n_dim: Divisor used to determine the number of convolved channels.
+                 The layer will convolve `C // n_dim` channels and leave the
+                 remaining channels untouched. Must be > 0, and
+                 `C // n_dim` must be >= 1 for a meaningful layer.
+        - key: PRNG key used to initialize the underlying convolution weights.
+        - kernel_size: Convolution kernel size (passed to `eqx.nn.Conv2d`).
+        - padding: Convolution padding (passed to `eqx.nn.Conv2d`). Use
+                   `"SAME"` to preserve spatial dimensions with `stride=1`.
+                   Integer or other forms supported by `eqx.nn.Conv2d` are also
+                   accepted, but must preserve H and W for the slice update.
+        - use_bias: Whether to include a bias term in the underlying convolution.
+        - **kwargs: Forwarded to `eqx.nn.Conv2d` (e.g., `dilation`, `groups`).
+        """
+        self.dim = in_channels // n_dim
+        assert self.dim >= 1, "in_channels // n_dim must be >= 1"
+        assert self.dim <= in_channels, (
+            "Computed convolved channels exceed total channels"
+        )
+
+        if isinstance(padding, str):
+            assert padding.upper() == "SAME", (
+                'When padding is a string, it must be "SAME"'
+            )
+        else:
+            # If padding is numeric, ensure it preserves H, W for stride=1
+            # For Conv2d with dilation d and kernel k, effective kernel is k_eff = (k - 1) * d + 1.
+            # Output H' = H + 2*pad - k_eff + 1; preserving H requires k_eff odd and pad = k_eff // 2.
+            dilation = kwargs.get("dilation", 1)
+            if isinstance(dilation, (tuple, list)):
+                # Require isotropic dilation for simplicity
+                assert len(dilation) == 2 and dilation[0] == dilation[1], (
+                    "dilation must be an int or an equal pair"
+                )
+                dilation = dilation[0]
+            assert isinstance(dilation, int) and dilation >= 1, (
+                "dilation must be a positive int"
+            )
+            assert isinstance(kernel_size, int) and kernel_size >= 1, (
+                "kernel_size must be a positive int"
+            )
+
+            k_eff = (kernel_size - 1) * dilation + 1
+            assert isinstance(padding, int) and padding >= 0, (
+                "padding must be a non-negative int"
+            )
+            assert k_eff % 2 == 1 and padding == k_eff // 2, (
+                "Integer padding must preserve spatial size: require effective kernel odd and "
+                f"padding == ((kernel_size-1)*dilation+1)//2; got k_eff={k_eff}, padding={padding}"
+            )
+
+        self.conv = eqx.nn.Conv2d(
+            in_channels=self.dim,
+            out_channels=self.dim,
+            kernel_size=kernel_size,
+            stride=1,
+            padding=padding,
+            use_bias=use_bias,
+            key=key,
+            **kwargs,
+        )
+
+    def __call__(
+        self,
+        x: Float[Array, "channels height width"],
+        *args,
+        **kwargs,
+    ) -> Float[Array, "channels height width"]:
+        c = self.dim
+        x1 = x[:c, :, :]
+        y1 = self.conv(x1)
+
+        return x.at[:c, :, :].set(y1)
 
-    This function recursively traverses the model and for each GenericGhostModule:
-    1. Fuses the parallel convolutional branches.
-    2. Replaces the training-only branches with Identity layers.
+
+class FasterNetBlock(eqx.Module):
+    """
+    FasterNet-style residual block with Partial Convolution-based spatial mixing
+    and pointwise MLP, adapted to Equinox/JAX.
+
+    Structure
+    - Spatial mixing: `PartialConv2d` applies a 3×3 convolution to the first
+      `C // n_dim` channels and leaves the remaining channels unchanged. This
+      reduces compute while keeping the tensor shape intact. See [1].
+    - Channel MLP: two pointwise (1*1) convolutions expand and then project
+      channels (`C -> mlp_ratio*C -> C`), with normalization and activation
+      in between.
+    - Regularization: includes dropout after the MLP and optional stochastic
+      depth (DropPath) on the residual branch.
+    - Residual: optional residual connection `y = x + DropPath(MLP(SpatialMix(x)))`.
+
+    Shape invariants
+    - Input: `[channels, height, width]`
+    - Output: `[channels, height, width]` (same spatial and channel dimensions)
+
+    References
+    - [1] Chen et al., "Run, Don't Walk: Chasing Higher FLOPS for Faster Neural
+          Networks" [arXiv:2303.03667](https://arxiv.org/abs/2303.03667).
+
+    Attributes
+    - residual: Whether to use a residual connection with stochastic depth.
+    - spatial_mixing: `PartialConv2d` performing partial 3×3 spatial mixing.
+    - pw_conv1: Pointwise convolution expanding channels to `mlp_ratio * C`.
+    - pw_conv2: Pointwise convolution projecting channels back to `C`.
+    - norm: Normalization layer applied on the expanded channels.
+    - act: Activation applied after normalization (defaults to identity if none).
+    - dropout: Dropout applied after the MLP.
+    - drop_path: Stochastic depth module for residual addition.
     """
 
-    def _finalize_leaf(module: GenericGhostModule) -> GenericGhostModule:
-        if not isinstance(module, GenericGhostModule):
-            return module
+    residual: bool = eqx.field(static=True)
 
-        fused_module = update_fused_ghostmodule(module)
+    spatial_mixing: eqx.nn.Conv2d
+    pw_conv1: eqx.nn.Conv2d
+    pw_conv2: eqx.nn.Conv2d
+    norm: eqx.Module
+    act: eqx.Module
+    dropout: eqx.nn.Dropout
+    drop_path: DropPathAdd
 
-        # identity_list = [eqx.nn.Identity()] * len(module.primary_rpr_conv)
-        final_module = eqx.tree_at(lambda m: m.primary_rpr_conv, fused_module, list())
-        final_module = eqx.tree_at(lambda m: m.cheap_rpr_conv, final_module, list())
-        final_module = eqx.tree_at(
-            lambda m: m.cheap_rpr_scale, final_module, eqx.nn.Identity()
+    def __init__(
+        self,
+        in_channels: int,
+        *,
+        key: PRNGKeyArray,
+        n_dim: int = 4,
+        mlp_ratio: int = 3,
+        kernel_size: int = 3,
+        padding: str | int = "SAME",
+        norm_layer: eqx.Module | None = eqx.nn.GroupNorm,
+        norm_max_group: int = 32,
+        act_layer: Callable | None = jax.nn.relu,
+        dropout: float = 0.0,
+        drop_path: float = 0.0,
+        norm_kwargs: dict = {},
+        residual: bool = True,
+        **kwargs,
+    ):
+        """
+        Initialize a FasterNetBlock.
+
+        Parameters
+        - in_channels: Number of input/output channels `C`.
+        - key: PRNG key used to initialize submodules. Internally split for
+          spatial mixing and pointwise convolutions.
+        - n_dim: Divisor determining the fraction of channels convolved by
+          `PartialConv2d`; the convolved channels are `C // n_dim`.
+        - mlp_ratio: Expansion ratio for the MLP (1×1 convs): hidden size is
+          `mlp_ratio * C`.
+        - kernel_size: Kernel size for the spatial mixing convolution (passed
+          to `PartialConv2d`).
+        - padding: Padding for the spatial mixing convolution. Use `"SAME"`
+          to preserve spatial dimensions.
+        - norm_layer: Normalization constructor applied on the expanded
+          channels. If `eqx.nn.GroupNorm`, the number of groups is chosen as
+          the largest power-of-two divisor of `hidden_channels` not exceeding
+          `norm_max_group`. If `None`, uses identity.
+        - norm_max_group: Maximum group count when using `GroupNorm`.
+        - act_layer: Callable used to construct an activation function. If
+          `None`, uses identity. Passed to `eqx.nn.Lambda`.
+        - dropout: Dropout probability applied after the MLP branch.
+        - drop_path: Stochastic depth probability on the residual branch.
+        - norm_kwargs: Extra keyword arguments forwarded to the normalization
+          layer constructor.
+        - residual: Whether to add the residual connection (with DropPath).
+        - **kwargs: Reserved for future extensions; forwarded where applicable.
+
+        Notes
+        - The block preserves `[H, W]`. Ensure `PartialConv2d` is configured
+          to preserve spatial dimensions (e.g., `"SAME"` padding).
+        - When `norm_layer` is not `GroupNorm`, it should accept the channel
+          count as its first argument.
+        - The same input/output channel count `C` is used throughout the block.
+        """
+
+        key_sm, key_pw1, key_pw2 = jr.split(key, 3)
+        self.residual = residual
+
+        self.spatial_mixing = PartialConv2d(
+            in_channels=in_channels, n_dim=n_dim, key=key_sm
         )
-        final_module = eqx.tree_at(
-            lambda m: m.cheap_rpr_skip, final_module, eqx.nn.Identity()
+
+        hidden_channels = mlp_ratio * in_channels
+        self.pw_conv1 = eqx.nn.Conv2d(
+            in_channels,
+            hidden_channels,
+            kernel_size=1,
+            stride=1,
+            padding="SAME",
+            use_bias=False,
+            key=key_pw1,
+        )
+        self.pw_conv2 = eqx.nn.Conv2d(
+            hidden_channels,
+            in_channels,
+            kernel_size=1,
+            stride=1,
+            padding="SAME",
+            use_bias=False,
+            key=key_pw2,
+        )
+
+        if norm_layer is not None:
+            if norm_layer == eqx.nn.GroupNorm:
+                num_groups = nearest_power_of_2_divisor(hidden_channels, norm_max_group)
+                self.norm = eqx.nn.GroupNorm(num_groups, hidden_channels, **norm_kwargs)
+            else:
+                self.norm = norm_layer(hidden_channels, **norm_kwargs)
+        else:
+            self.norm = eqx.nn.Identity()
+
+        self.dropout = eqx.nn.Dropout(dropout)
+        self.drop_path = DropPathAdd(drop_path)
+        self.act = eqx.nn.Lambda(act_layer) if act_layer else eqx.nn.Identity()
+
+    def __call__(
+        self,
+        x: Float[Array, "channels height width"],
+        key: PRNGKeyArray,
+        inference: Optional[bool] = None,
+    ) -> Float[Array, "channels height width"]:
+        key_dropout, key_droppath = jr.split(key, 2)
+        x1 = self.spatial_mixing(x)
+        out = self.dropout(
+            self.pw_conv2(self.act(self.norm(self.pw_conv1(x1)))),
+            inference=inference,
+            key=key_dropout,
         )
 
-        return final_module
+        if self.residual:
+            out = self.drop_path(x, out, inference=inference, key=key_droppath)
 
-    is_g_module = lambda m: isinstance(m, GenericGhostModule)
-    return jax.tree_util.tree_map(_finalize_leaf, model, is_leaf=is_g_module)
+        return out

equimo-0.4.3/src/equimo/__init__.py

@@ -1 +0,0 @@
-__version__ = "0.4.3"