Equimo 0.4.1__tar.gz → 0.4.3__tar.gz

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

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

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

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

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

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

equimo-0.4.3/src/equimo/__init__.py

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

{equimo-0.4.1 → equimo-0.4.3}/src/equimo/layers/attention.py

@@ -1766,6 +1766,8 @@ class RFAttention(eqx.Module):
 
 
 class RFAttentionBlock(eqx.Module):
+    residual: bool = eqx.field(static=True)
+
     context_module: RFAttention
     local_module: MBConv
     drop_path1: DropPathAdd
@@ -1789,9 +1791,11 @@ class RFAttentionBlock(eqx.Module):
         local_drop: float = 0.0,
         drop_path: float | List[float] = 0.0,
         residual_mbconv: bool = False,
+        residual: bool = True,
         **kwargs,
     ):
         key_context, key_local = jr.split(key, 2)
+        self.residual = residual
 
         if isinstance(drop_path, list):
             if len(drop_path) != 2:
@@ -1841,20 +1845,15 @@ class RFAttentionBlock(eqx.Module):
         key_context, key_local, key_dr1, key_dr2 = jr.split(key, 4)
 
         # TODO: some prenorm?
-        x = self.drop_path1(
-            x,
-            self.context_module(x, inference=inference, key=key_context),
-            inference=inference,
-            key=key_dr1,
-        )
-        x = self.drop_path2(
-            x,
-            self.local_module(x, inference=inference, key=key_local),
-            inference=inference,
-            key=key_dr2,
-        )
+        x1 = self.context_module(x, inference=inference, key=key_context)
+        if self.residual:
+            x1 = self.drop_path1(x, x1, inference=inference, key=key_dr1)
 
-        return x
+        x2 = self.local_module(x, inference=inference, key=key_local)
+        if self.residual:
+            x2 = self.drop_path2(x1, x2, inference=inference, key=key_dr2)
+
+        return x2
 
 
 class ConvAttention(eqx.Module):

{equimo-0.4.1 → equimo-0.4.3}/src/equimo/models/reduceformer.py

@@ -159,7 +159,7 @@ class ReduceFormer(eqx.Module):
         block_types: list[Literal["conv", "attention"]],
         *,
         key: PRNGKeyArray,
-        heads_dim: int = 32,
+        head_dim: int = 32,
         expand_ratio: float = 4.0,
         norm_layer: eqx.Module | str = eqx.nn.GroupNorm,
         act_layer: Callable | str = jax.nn.hard_swish,
@@ -227,6 +227,7 @@ class ReduceFormer(eqx.Module):
                 depth=depth,
                 block_type=block_type,
                 stride=2,
+                head_dim=head_dim,
                 expand_ratio=expand_ratio,
                 norm_layer=norm_layer,
                 act_layer=act_layer,
@@ -335,7 +336,7 @@ def reduceformer_backbone_b1(**kwargs) -> ReduceFormer:
             "attention",
             "attention",
         ],
-        heads_dim=16,
+        head_dim=16,
         **kwargs,
     )
     return backbone
@@ -352,7 +353,7 @@ def reduceformer_backbone_b2(**kwargs) -> ReduceFormer:
             "attention",
             "attention",
         ],
-        heads_dim=32,
+        head_dim=32,
         **kwargs,
     )
     return backbone
@@ -369,7 +370,7 @@ def reduceformer_backbone_b3(**kwargs) -> ReduceFormer:
             "attention",
             "attention",
         ],
-        heads_dim=32,
+        head_dim=32,
         **kwargs,
     )
     return backbone

equimo-0.4.1/src/equimo/__init__.py

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