Equimo 0.4.7__tar.gz → 0.5.0a2__tar.gz

{equimo-0.4.7/src/Equimo.egg-info → equimo-0.5.0a2}/PKG-INFO

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

{equimo-0.4.7 → equimo-0.5.0a2}/pyproject.toml

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

{equimo-0.4.7 → equimo-0.5.0a2/src/Equimo.egg-info}/PKG-INFO

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

{equimo-0.4.7 → equimo-0.5.0a2}/src/Equimo.egg-info/SOURCES.txt

@@ -29,6 +29,7 @@ src/equimo/layers/sharing.py
 src/equimo/layers/squeeze_excite.py
 src/equimo/layers/wavelet.py
 src/equimo/models/__init__.py
+src/equimo/models/attnet.py
 src/equimo/models/emamodel.py
 src/equimo/models/fastervit.py
 src/equimo/models/lowformer.py

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

@@ -0,0 +1 @@
+__version__ = "0.5.0-alpha.2"

{equimo-0.4.7 → equimo-0.5.0a2}/src/equimo/layers/attention.py

@@ -356,8 +356,8 @@ class AttentionBlock(eqx.Module):
         self.norm = norm_layer(dim, eps=eps)
 
         if init_values:
-            self.ls1 = LayerScale(dim, init_values=init_values)
-            self.ls2 = LayerScale(dim, init_values=init_values)
+            self.ls1 = LayerScale(dim, axis=1, init_values=init_values)
+            self.ls2 = LayerScale(dim, axis=1, init_values=init_values)
         else:
             self.ls1 = self.ls2 = eqx.nn.Identity()
 
@@ -538,8 +538,8 @@ class HATBlock(eqx.Module):
         self.norm2 = norm_layer(dim, eps=eps)
 
         if init_values:
-            self.ls1 = LayerScale(dim, init_values=init_values)
-            self.ls2 = LayerScale(dim, init_values=init_values)
+            self.ls1 = LayerScale(dim, axis=1, init_values=init_values)
+            self.ls2 = LayerScale(dim, axis=1, init_values=init_values)
         else:
             self.ls1 = self.ls2 = eqx.nn.Identity()
 
@@ -611,17 +611,17 @@ class HATBlock(eqx.Module):
                 dim, rank=2, seq_len=cr_tokens_total, key=key_hatposemb
             )
             self.hat_ls1 = (
-                LayerScale(dim, init_values=init_values)
+                LayerScale(dim, axis=1, init_values=init_values)
                 if init_values
                 else eqx.nn.Identity()
             )
             self.hat_ls2 = (
-                LayerScale(dim, init_values=init_values)
+                LayerScale(dim, axis=1, init_values=init_values)
                 if init_values
                 else eqx.nn.Identity()
             )
             self.hat_ls3 = (
-                LayerScale(dim, init_values=init_values)
+                LayerScale(dim, axis=1, init_values=init_values)
                 if init_values
                 else eqx.nn.Identity()
             )
@@ -1425,8 +1425,8 @@ class PartialFormerBlock(eqx.Module):
         self.norm2 = norm_layer(dim, eps=eps)
 
         if init_values:
-            self.ls1 = LayerScale(dim, init_values=init_values)
-            self.ls2 = LayerScale(dim, init_values=init_values)
+            self.ls1 = LayerScale(dim, axis=1, init_values=init_values)
+            self.ls2 = LayerScale(dim, axis=1, init_values=init_values)
         else:
             self.ls1 = self.ls2 = eqx.nn.Identity()
 

{equimo-0.4.7 → equimo-0.5.0a2}/src/equimo/layers/convolution.py

@@ -39,7 +39,7 @@ class ConvBlock(eqx.Module):
     norm2: eqx.Module
     drop_path1: DropPathAdd
     act: Callable
-    ls1: LayerScale | None
+    ls1: LayerScale | eqx.nn.Identity
 
     def __init__(
         self,
@@ -50,7 +50,7 @@ class ConvBlock(eqx.Module):
         kernel_size: int = 3,
         stride: int = 1,
         padding: int = 1,
-        act_layer: Callable = jax.nn.gelu,
+        act_layer: Callable | None = jax.nn.gelu,
         norm_max_group: int = 32,
         drop_path: float = 0.0,
         init_values: float | None = None,
@@ -87,7 +87,7 @@ class ConvBlock(eqx.Module):
             key=key_conv1,
         )
         self.norm1 = eqx.nn.GroupNorm(num_groups1, hidden_dim)
-        self.act = act_layer
+        self.act = act_layer if act_layer is not None else eqx.nn.Identity()
         self.conv2 = eqx.nn.Conv(
             num_spatial_dims=2,
             in_channels=hidden_dim,
@@ -103,18 +103,11 @@ class ConvBlock(eqx.Module):
         dpr = drop_path[0] if isinstance(drop_path, list) else float(drop_path)
         self.drop_path1 = DropPathAdd(dpr)
 
-        self.ls1 = LayerScale(dim, init_values=init_values) if init_values else None
-
-    def permute(
-        self, x: Float[Array, "channels height width"]
-    ) -> Float[Array, "height width channels"]:
-        return rearrange(x, "c h w -> h w c")
-
-    def depermute(
-        self,
-        x: Float[Array, "height width channels"],
-    ) -> Float[Array, "channels height width"]:
-        return rearrange(x, "h w c -> c h w")
+        self.ls1 = (
+            LayerScale(dim, init_values=init_values)
+            if init_values
+            else eqx.nn.Identity()
+        )
 
     def __call__(
         self,
@@ -124,8 +117,7 @@ class ConvBlock(eqx.Module):
     ) -> Float[Array, "channels height width"]:
         x2 = self.act(self.norm1(self.conv1(x)))
         x2 = self.norm2(self.conv2(x2))
-        if self.ls1 is not None:
-            x2 = self.depermute(jax.vmap(jax.vmap(self.ls1))(self.permute(x2)))
+        x2 = self.ls1(x2)
 
         return self.drop_path1(x, x2, inference=inference, key=key)
 
@@ -2017,3 +2009,309 @@ class FasterNetBlock(eqx.Module):
             out = self.drop_path(x, out, inference=inference, key=key_droppath)
 
         return out
+
+
+class GLUConv(eqx.Module):
+    conv1: eqx.nn.Conv2d
+    conv2: eqx.nn.Conv2d
+    dwconv: eqx.nn.Conv2d | eqx.nn.Identity
+    norm: eqx.nn.GroupNorm | eqx.nn.Identity
+    act: Callable
+    dr1: eqx.nn.Dropout
+    dr2: eqx.nn.Dropout
+
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_channels: int,
+        *,
+        glu_norm: bool = True,
+        glu_dwconv: bool = False,
+        act_layer: Callable = jax.nn.gelu,
+        dropout: float = 0.0,
+        key: PRNGKeyArray,
+    ):
+        key_conv1, key_conv2, key_dwconv = jr.split(key, 3)
+
+        self.conv1 = eqx.nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=hidden_channels * 2,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            use_bias=True,
+            key=key_conv1,
+        )
+        self.conv2 = eqx.nn.Conv2d(
+            in_channels=hidden_channels,
+            out_channels=in_channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            use_bias=True,
+            key=key_conv2,
+        )
+        self.dwconv = (
+            eqx.nn.Conv2d(
+                in_channels=hidden_channels,
+                out_channels=hidden_channels,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                groups=hidden_channels,
+                use_bias=True,
+                key=key_dwconv,
+            )
+            if glu_dwconv
+            else eqx.nn.Identity()
+        )
+
+        num_groups = nearest_power_of_2_divisor(in_channels, 32)
+        self.norm = (
+            eqx.nn.GroupNorm(num_groups, hidden_channels)
+            if glu_norm
+            else eqx.nn.Identity()
+        )
+
+        self.act = act_layer
+        self.dr1 = eqx.nn.Dropout(dropout)
+        self.dr2 = eqx.nn.Dropout(dropout)
+
+    def __call__(
+        self,
+        x: Float[Array, "channels height width"],
+        key: PRNGKeyArray,
+        inference: Optional[bool] = None,
+    ) -> Float[Array, "channels height width"]:
+        key_dr1, key_dr2 = jr.split(key, 2)
+
+        x, v = jnp.split(self.conv1(x), 2)
+
+        x = self.norm(self.act(self.dwconv(x)) * v)
+        x = self.dr1(x, inference=inference, key=key_dr1)
+        x = self.conv2(x)
+        x = self.dr2(x, inference=inference, key=key_dr2)
+
+        return x
+
+
+class ATConv(eqx.Module):
+    """
+    Attentive Convolution (ATConv) layer that generates per-sample, per-channel 3×3 dynamic kernels
+    and applies them as a depthwise grouped 2D convolution with linear-time cost, capturing adaptive
+    routing and lateral inhibition principles identified as key advantages of self-attention over static
+    convolutions[1].
+
+    The operator first projects features, synthesizes a per-channel HWIO kernel from pooled spatial
+    summaries, subtracts a learnable fraction of its channel-wise mean to induce inhibition, then
+    performs a depthwise convolution followed by a 1×1 projection, matching the reference design
+    that unifies attention-like expressivity with convolutional efficiency[1].
+
+    Referencea:
+        [1]. Yu, et al., Attentive Convolution: Unifying the Expressivity of Self-Attention with Convolutional
+             Efficiency. 2025. arXiv:2510.20092
+    """
+
+    kernel_size: int = eqx.field(static=True)
+    padding: int = eqx.field(static=True)
+
+    x_proj: eqx.nn.Conv2d
+    proj: eqx.nn.Conv2d
+    kernel_proj: eqx.nn.Conv2d
+    kernel_act: Callable
+    kernel_gen: eqx.nn.Linear
+    pool: eqx.nn.AdaptiveAvgPool1d
+    difference_control: jax.Array
+
+    def __init__(
+        self,
+        in_channels: int,
+        *,
+        kernel_size: int = 3,
+        act_layer: Callable = jax.nn.gelu,
+        use_bias: bool = True,
+        key: PRNGKeyArray,
+    ):
+        key_xproj, key_proj, key_kp, key_kg = jr.split(key, 4)
+
+        self.kernel_size = kernel_size
+        k2 = kernel_size**2
+        self.padding = kernel_size // 2
+
+        self.x_proj = eqx.nn.Conv2d(
+            in_channels,
+            in_channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            use_bias=use_bias,
+            key=key_xproj,
+        )
+        self.proj = eqx.nn.Conv2d(
+            in_channels,
+            in_channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            use_bias=use_bias,
+            key=key_proj,
+        )
+        self.kernel_proj = eqx.nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=in_channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            use_bias=use_bias,
+            key=key_kp,
+        )
+        self.kernel_act = act_layer
+        self.kernel_gen = eqx.nn.Linear(k2, k2, use_bias=use_bias, key=key_kg)
+        self.pool = eqx.nn.AdaptiveAvgPool1d(target_shape=k2)
+        self.difference_control = jnp.zeros((1, 1, 1, in_channels))
+
+    def _generate_kernels(
+        self, x: Float[Array, "channels height width"]
+    ) -> Float[Array, "k k 1 channels"]:
+        kernels = self.kernel_proj(x)
+        kernels = self.pool(rearrange(kernels, "c h w -> c (h w)"))
+        kernels = self.kernel_act(kernels)
+        kernels = jax.vmap(self.kernel_gen)(kernels)
+        kernels = rearrange(
+            kernels, "c (k1 k2) -> k1 k2 1 c", k1=self.kernel_size, k2=self.kernel_size
+        )
+
+        return kernels
+
+    def _apply_kernel_difference(
+        self, kernels: Float[Array, "channels k k"]
+    ) -> Float[Array, "channels k k"]:
+        mean_kernels = kernels.mean(axis=(0, 1, 2), keepdims=True)
+        factor = jax.nn.sigmoid(self.difference_control)
+
+        return kernels - factor * mean_kernels
+
+    def __call__(
+        self,
+        x: Float[Array, "channels height width"],
+        key: PRNGKeyArray,
+        inference: Optional[bool] = None,
+    ) -> Float[Array, "channels height width"]:
+        kernels = self._generate_kernels(x)
+        kernels = self._apply_kernel_difference(kernels)
+
+        x = self.x_proj(x)
+        C, H, W = x.shape
+        x = jax.lax.conv_general_dilated(
+            x[None, ...],  # (N, C, H, W)
+            kernels,  # (H, W, I=1, O=C)
+            window_strides=(1, 1),
+            padding=((self.padding, self.padding), (self.padding, self.padding)),
+            lhs_dilation=None,
+            rhs_dilation=None,
+            dimension_numbers=("NCHW", "HWIO", "NCHW"),
+            feature_group_count=C,  # depthwise
+        )[0]
+
+        x = self.proj(x)
+
+        return x
+
+
+class ATConvBlock(eqx.Module):
+    token_mixer: ATConv
+    channel_mixer: GLUConv
+    ls1: LayerScale | eqx.nn.Identity
+    ls2: LayerScale | eqx.nn.Identity
+    norm1: eqx.nn.LayerNorm
+    norm2: eqx.nn.LayerNorm
+    drop_path1: DropPathAdd
+    drop_path2: DropPathAdd
+
+    def __init__(
+        self,
+        in_channels: int,
+        *,
+        kernel_size: int = 3,
+        exp_rate: float = 4.0,
+        act_layer: Callable = jax.nn.gelu,
+        glu_norm: bool = True,
+        glu_dwconv: bool = False,
+        use_bias: bool = True,
+        dropout: float = 0.0,
+        drop_path: float | list[float] = 0.0,
+        use_layer_scale: bool = True,
+        key: PRNGKeyArray,
+    ):
+        key_tm, key_cm = jr.split(key, 2)
+
+        hidden_features = int(in_channels * exp_rate)
+        # NOTE: slightly different from the original paper
+        glu_hidden_features = 32 * round(int(2 * hidden_features / 3) / 32)
+
+        self.token_mixer = ATConv(
+            in_channels,
+            kernel_size=kernel_size,
+            act_layer=act_layer,
+            use_bias=use_bias,
+            key=key_tm,
+        )
+        self.channel_mixer = GLUConv(
+            in_channels,
+            hidden_channels=glu_hidden_features,
+            glu_norm=glu_norm,
+            glu_dwconv=glu_dwconv,
+            act_layer=act_layer,
+            dropout=dropout,
+            key=key_cm,
+        )
+
+        self.ls1 = LayerScale(in_channels) if use_layer_scale else eqx.nn.Identity()
+        self.ls2 = LayerScale(in_channels) if use_layer_scale else eqx.nn.Identity()
+
+        self.norm1 = eqx.nn.LayerNorm(in_channels)
+        self.norm2 = eqx.nn.LayerNorm(in_channels)
+
+        if isinstance(drop_path, list):
+            if len(drop_path) != 2:
+                raise AssertionError(
+                    f"`drop_path` needs to have 2 elements, got {len(drop_path)} ({drop_path})."
+                )
+            dr1, dr2 = drop_path
+            dr1 = float(dr1)
+            dr2 = float(dr2)
+        else:
+            dr1 = dr2 = float(drop_path)
+
+        self.drop_path1 = DropPathAdd(dr1)
+        self.drop_path2 = DropPathAdd(dr2)
+
+    def __call__(
+        self,
+        x: Float[Array, "channels height width"],
+        key: PRNGKeyArray,
+        inference: Optional[bool] = None,
+    ) -> Float[Array, "channels height width"]:
+        key_tm, key_cm, key_dr1, key_dr2 = jr.split(key, 4)
+
+        norm1 = jax.vmap(
+            jax.vmap(self.norm1, in_axes=1, out_axes=1), in_axes=1, out_axes=1
+        )
+        norm2 = jax.vmap(
+            jax.vmap(self.norm2, in_axes=1, out_axes=1), in_axes=1, out_axes=1
+        )
+
+        x = self.drop_path1(
+            x,
+            self.ls1(self.token_mixer(norm1(x), inference=inference, key=key_tm)),
+            inference=inference,
+            key=key_dr1,
+        )
+        x = self.drop_path2(
+            x,
+            self.ls2(self.channel_mixer(norm2(x), inference=inference, key=key_cm)),
+            inference=inference,
+            key=key_dr2,
+        )
+
+        return x

equimo-0.5.0a2/src/equimo/layers/downsample.py

@@ -0,0 +1,209 @@
+from typing import Callable, Literal, Optional
+
+import equinox as eqx
+import jax
+import jax.random as jr
+from jaxtyping import Array, Float, PRNGKeyArray
+
+from equimo.layers.convolution import ConvBlock, SingleConvBlock
+from equimo.layers.generic import Residual
+from equimo.layers.patch import SEPatchMerging
+from equimo.utils import nearest_power_of_2_divisor
+
+
+class ConvNormDownsampler(eqx.Module):
+    """A module that performs spatial downsampling using strided convolution.
+
+    This module reduces spatial dimensions (height and width) by a factor of 2
+    while optionally increasing the channel dimension. Uses a 3x3 strided
+    convolution for downsampling.
+
+    Attributes:
+        reduction: Convolutional layer that performs the downsampling
+    """
+
+    downsampler: eqx.nn.Sequential
+
+    def __init__(
+        self,
+        in_channels: int,
+        *,
+        out_channels: int | None = None,
+        act_layer: Callable = jax.nn.gelu,
+        use_bias: bool = False,
+        use_norm: bool = True,
+        mode: Literal["double", "simple"] = "simple",
+        key: PRNGKeyArray,
+    ):
+        out_channels = out_channels if out_channels else 2 * in_channels
+
+        match mode:
+            case "simple":
+                self.downsampler = eqx.nn.Sequential(
+                    [
+                        eqx.nn.Conv2d(
+                            in_channels=in_channels,
+                            out_channels=out_channels,
+                            kernel_size=3,
+                            stride=2,
+                            padding=1,
+                            use_bias=False,
+                            key=key,
+                        ),
+                        eqx.nn.GroupNorm(
+                            nearest_power_of_2_divisor(out_channels, 32), out_channels
+                        )
+                        if use_norm
+                        else eqx.nn.Identity(),
+                    ]
+                )
+            case "double":
+                self.downsampler = eqx.nn.Sequential(
+                    [
+                        eqx.nn.Conv2d(
+                            in_channels=in_channels,
+                            out_channels=(_d := out_channels // 2),
+                            kernel_size=3,
+                            stride=2,
+                            padding=1,
+                            use_bias=use_bias,
+                            key=jr.fold_in(key, 0),
+                        ),
+                        eqx.nn.GroupNorm(nearest_power_of_2_divisor(_d, 32), _d)
+                        if use_norm
+                        else eqx.nn.Identity(),
+                        eqx.nn.Lambda(act_layer),
+                        eqx.nn.Conv2d(
+                            in_channels=_d,
+                            out_channels=out_channels,
+                            kernel_size=3,
+                            stride=2,
+                            padding=1,
+                            use_bias=use_bias,
+                            key=jr.fold_in(key, 1),
+                        ),
+                        eqx.nn.GroupNorm(
+                            nearest_power_of_2_divisor(out_channels, 32), out_channels
+                        )
+                        if use_norm
+                        else eqx.nn.Identity(),
+                    ]
+                )
+
+    def __call__(self, x, *args, **kwargs):
+        return self.downsampler(x)
+
+
+class PWSEDownsampler(eqx.Module):
+    """Downsampling module for spatial feature reduction.
+
+    Combines convolution blocks and patch merging to reduce spatial dimensions
+    while increasing feature channels. Uses residual connections and alternates
+    between depthwise and pointwise convolutions.
+
+    Attributes:
+        conv1: First depthwise convolution with residual
+        conv2: First pointwise convolution block
+        conv3: Second depthwise convolution with residual
+        conv4: Second pointwise convolution block
+        patch_merging: Squeeze-and-excitation patch merging layer
+    """
+
+    conv1: eqx.Module
+    conv2: eqx.Module
+    conv3: eqx.Module
+    conv4: eqx.Module
+    patch_merging: eqx.Module
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        *,
+        key: PRNGKeyArray,
+        drop_path: float = 0.0,
+        **kwargs,
+    ):
+        key_conv1, key_conv2, key_conv3, key_conv4, key_pm = jr.split(key, 5)
+        self.conv1 = Residual(
+            SingleConvBlock(
+                in_channels,
+                in_channels,
+                act_layer=None,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                groups=in_channels,
+                key=key_conv1,
+            ),
+            drop_path=drop_path,
+        )
+        self.conv2 = ConvBlock(
+            in_channels,
+            hidden_in_channels=in_channels * 2,
+            act_layer=None,
+            drop_path=drop_path,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            key=key_conv2,
+        )
+        self.patch_merging = SEPatchMerging(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            key=key_pm,
+        )
+        self.conv3 = Residual(
+            SingleConvBlock(
+                out_channels,
+                out_channels,
+                act_layer=None,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+                groups=out_channels,
+                key=key_conv3,
+            ),
+            drop_path=drop_path,
+        )
+        self.conv4 = ConvBlock(
+            out_channels,
+            hidden_in_channels=out_channels * 2,
+            act_layer=None,
+            drop_path=drop_path,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            key=key_conv4,
+        )
+
+    def __call__(
+        self,
+        x: Float[Array, "seqlen dim"],
+        key: PRNGKeyArray,
+        inference: Optional[bool] = None,
+    ) -> Float[Array, "seqlen dim"]:
+        """Apply downsampling to input features.
+
+        Args:
+            x: Input feature tensor
+            inference: Whether to enable dropout during inference
+            key: PRNG key for random operations
+
+        Returns:
+            Downsampled feature tensor with increased channels
+        """
+        key_conv1, key_conv2, key_conv3, key_conv4 = jr.split(key, 4)
+        x = self.conv2(
+            self.conv1(x, inference=inference, key=key_conv1),
+            inference=inference,
+            key=key_conv2,
+        )
+        x = self.patch_merging(x)
+        x = self.conv4(
+            self.conv3(x, inference=inference, key=key_conv3),
+            inference=inference,
+            key=key_conv4,
+        )
+
+        return x