Equimo 0.4.2__tar.gz → 0.4.3__tar.gz

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

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

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

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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: Equimo
-Version: 0.4.2
+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.2 → 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.2/src/equimo/__init__.py

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