Equimo 0.4.6a2__tar.gz → 0.4.6a3__tar.gz

{equimo-0.4.6a2/src/Equimo.egg-info → equimo-0.4.6a3}/PKG-INFO

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

{equimo-0.4.6a2 → equimo-0.4.6a3}/pyproject.toml

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

{equimo-0.4.6a2 → equimo-0.4.6a3/src/Equimo.egg-info}/PKG-INFO

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

equimo-0.4.6a3/src/equimo/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.4.6a3"

{equimo-0.4.6a2 → equimo-0.4.6a3}/src/equimo/layers/convolution.py

@@ -1300,16 +1300,33 @@ class GenericGhostModule(eqx.Module):
         inference: Optional[bool] = None,
     ):
         use_inference = self.inference if inference is None else inference
-        out = self.inference_features(x) if use_inference else self.training_features(x)
+        use_inference = jnp.asarray(use_inference, dtype=bool)
 
-        if self.mode == "shortcut":
-            res = self.short_conv(self.pool2(x))
-            gating_signal = jax.image.resize(
+        out = jax.lax.cond(
+            use_inference,
+            lambda x_: self.inference_features(x_),
+            lambda x_: self.training_features(x_),
+            operand=x,
+        )
+
+        shortcut_pred = bool(self.mode == "shortcut")
+
+        def _shortcut_branch(ox):
+            out_, x_ = ox
+            res = self.short_conv(self.pool2(x_))
+            gating = jax.image.resize(
                 image=jax.nn.sigmoid(res / self.gate_scale),
-                shape=(res.shape[0], out.shape[1], out.shape[2]),
+                shape=(res.shape[0], out_.shape[1], out_.shape[2]),
                 method="nearest",
             )
-            out = out[: self.out_channels, :, :] * gating_signal
+            return out_[: self.out_channels, :, :] * gating
+
+        out = jax.lax.cond(
+            shortcut_pred,
+            _shortcut_branch,
+            lambda ox: ox[0],
+            operand=(out, x),
+        )
 
         return out
 
@@ -1347,7 +1364,6 @@ class GhostBottleneck(eqx.Module):
         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,
@@ -1490,30 +1506,41 @@ class GhostBottleneck(eqx.Module):
         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)
+        use_inference = jax.numpy.asarray(use_inference, dtype=bool)
 
+        k_g1, k_g2 = jr.split(key, 2)
         residual = x
 
         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 = jax.lax.cond(
+                use_inference,
+                lambda y: self.dw_shared_norm(self.dw_conv(y)),
+                lambda y: self.dw_shared_norm(
+                    jax.tree_util.tree_reduce(
+                        operator.add,
+                        (
+                            (
+                                []
+                                if isinstance(self.dw_rpr_scale, eqx.nn.Identity)
+                                else [self.dw_rpr_scale(y)]
+                            )
+                            + [conv(y) for conv in self.dw_rpr_conv]
+                        ),
+                    )
+                ),
+                operand=x,
+            )
 
         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)
+            x = x + self.shortcut(residual)
         elif self.allow_identity_residual:
-            x += residual
+            x = x + residual
 
         return x
 

{equimo-0.4.6a2 → equimo-0.4.6a3}/src/equimo/layers/wavelet.py

@@ -63,7 +63,7 @@ def haar_dwt_split(
 
 
 class HWDConv(eqx.Module):
-    mode: Literal["h_discard", "accurate"] = eqx.field(static=True)
+    mode: Literal["h_discard", "band_grouped", "accurate"] = eqx.field(static=True)
 
     pre_norm: eqx.nn.GroupNorm
     proj: eqx.nn.Conv2d
@@ -78,31 +78,43 @@ class HWDConv(eqx.Module):
         kernel_size: int = 1,
         use_bias: bool = False,
         act: Callable = jax.nn.relu,
-        mode: Literal["h_discard", "accurate"] = "accurate",
+        mode: Literal["h_discard", "band_grouped", "accurate"] = "accurate",
         dropout: float = 0.0,
         key: PRNGKeyArray,
     ):
         self.mode = mode
 
-        # Channels entering the projection (and thus the pre-norm)
         channels_in = in_channels if mode == "h_discard" else 4 * in_channels
 
-        # Choose groups so no group crosses sub-band boundaries.
-        # For accurate mode: total_groups = 4 * groups_per_band, each band has groups_per_band groups.
-        # For h_discard mode: total_groups = groups_per_band (single band).
+        # GroupNorm: align groups to sub-band boundaries
+        # groups_per_band divides in_channels; total_groups = 4 * groups_per_band for 4 bands
         groups_per_band = nearest_power_of_2_divisor(in_channels, 32)
         total_groups = groups_per_band if mode == "h_discard" else 4 * groups_per_band
-
         self.pre_norm = eqx.nn.GroupNorm(channels=channels_in, groups=total_groups)
 
+        # Projection 1x1 conv
+        # - accurate: dense projection (groups=1)
+        # - band_grouped: independent per-band projections (groups=4), no cross-band mixing
+        # - h_discard: single-band dense projection (groups=1)
+        proj_groups = 1
+        if mode == "band_grouped":
+            if out_channels % 4 != 0:
+                raise ValueError(
+                    "band_grouped mode requires out_channels to be divisible by 4 "
+                    "so each sub-band produces an equal share."
+                )
+            proj_groups = 4
+
         self.proj = eqx.nn.Conv2d(
-            channels_in,
-            out_channels,
+            in_channels=channels_in,
+            out_channels=out_channels,
             kernel_size=kernel_size,
             use_bias=use_bias,
             padding="SAME",
+            groups=proj_groups,
             key=key,
         )
+
         self.act = act
 
         self.dropout = eqx.nn.Dropout(dropout)
@@ -115,11 +127,13 @@ class HWDConv(eqx.Module):
         if self.mode == "h_discard":
             y = LL  # (C, H/2, W/2)
         else:
+            # Keep sub-bands contiguous and ordered so GN groups and grouped conv line up
             y = jnp.concatenate([LL, HL, LH, HH], axis=0)  # (4*C, H/2, W/2)
 
+        # Band-aligned normalization (no cross-band statistics)
         y = self.pre_norm(y)
 
-        y = self.proj(y)  # (out_ch, H/2, W/2)
+        y = self.proj(y)  # (out_channels, H/2, W/2)
         y = self.act(y)
 
         y = self.dropout(y, inference=inference, key=key)

equimo-0.4.6a2/src/equimo/__init__.py

@@ -1 +0,0 @@
-__version__ = "0.4.6a2"