Equimo 0.4.0a2__tar.gz → 0.4.0a3__tar.gz

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

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

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

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

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

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

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

@@ -0,0 +1 @@
+__version__ = "0.4.0-alpha.3"

{equimo-0.4.0a2 → equimo-0.4.0a3}/src/equimo/layers/attention.py

@@ -97,7 +97,7 @@ class Attention(eqx.Module):
             attn = jnp.where(mask == 0, jnp.finfo(attn.dtype).min, attn)
 
         attn = jax.nn.softmax(attn, axis=-1)
-        attn = self.attn_drop(attn, inference=inference, key=key1)
+        attn = self.attn_drop(attn, inference=infeoence, key=key1)
 
         x = jnp.einsum("hqk,hkd->hqd", attn, v)
         x = rearrange(x, "h s d -> s (h d)")
@@ -1591,6 +1591,7 @@ class RFAttention(eqx.Module):
     qkv: eqx.nn.Conv2d
     aggreg: list[eqx.nn.Conv2d]
     proj: SingleConvBlock
+    proj_drop: eqx.nn.Dropout
 
     def __init__(
         self,
@@ -1604,6 +1605,7 @@ class RFAttention(eqx.Module):
         scales: Sequence[int] = (5,),
         use_bias: bool = False,
         kernel_func: Callable = jax.nn.relu,
+        proj_drop: float = 0.0,
         # TODO: Benchmark against LN, RMSN, NsLN
         norm_layer: eqx.Module = eqx.nn.GroupNorm,
         norm_kwargs: dict = {},
@@ -1649,12 +1651,14 @@ class RFAttention(eqx.Module):
             key=key_proj,
         )
 
+        self.proj_drop = eqx.nn.Dropout(proj_drop)
+
     def __call__(
         self,
-        x: Float[Array, "seqlen height width"],
+        x: Float[Array, "dim height width"],
         key: PRNGKeyArray,
         inference: Optional[bool] = None,
-    ) -> Float[Array, "seqlen height width"]:
+    ) -> Float[Array, "dim height width"]:
         qkv_base = self.qkv(x)
 
         aggregated_qkvs = [op(qkv_base) for op in self.aggreg]
@@ -1671,12 +1675,13 @@ class RFAttention(eqx.Module):
         k = self.kernel_func(k)
 
         sum_k = jnp.sum(k, axis=(-1, -2), keepdims=True)
-        sum_v = jnp.sum(v, axis=(-1, -2), keepdims=True)
+        sum_v = jnp.sum(v * sum_k, axis=(-1, -2), keepdims=True)
         sum_kv = jnp.sum(k * sum_v, axis=(-1, -2), keepdims=True)
         sum_q = jnp.sum(q, axis=0, keepdims=True)
 
         out = (q * sum_kv) / (sum_q * sum_k + self.eps)
         out = self.proj(out)
+        out = self.proj_drop(out, inference=inference, key=key)
 
         return out
 

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

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