Equimo 0.3.0__tar.gz → 0.3.2__tar.gz

{equimo-0.3.0 → equimo-0.3.2}/Equimo.egg-info/PKG-INFO

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

{equimo-0.3.0 → equimo-0.3.2}/PKG-INFO

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

equimo-0.3.2/equimo/__init__.py

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

{equimo-0.3.0 → equimo-0.3.2}/equimo/layers/norm.py

@@ -31,6 +31,8 @@ class RMSNormGated(eqx.Module):
         self,
         x: Float[Array, "dim"],
         z: Optional[Float[Array, "dim"]] = None,
+        *args,
+        **kwargs,
     ) -> Float[Array, "dim"]:
         """Apply RMS normalization with optional gating.
 
@@ -58,12 +60,10 @@ class LayerScale(eqx.Module):
     dampening the contribution of each layer.
 
     Attributes:
-        init_values: Initial scale value (static)
         gamma: Learnable scale parameters of size dim
     """
 
-    init_values: float = eqx.field(static=True)
-    gamma: Optional[Float[Array, "dim"]]
+    gamma: Float[Array, "dim"]
 
     def __init__(self, dim: int, init_values: float):
         """Initialize LayerScale.
@@ -72,8 +72,7 @@ class LayerScale(eqx.Module):
             dim: Dimension of the input features
             init_values: Initial value for all scaling factors
         """
-        self.init_values = init_values
-        self.gamma = jnp.repeat(self.init_values, dim)
+        self.gamma = jnp.repeat(init_values, dim)
 
     def __call__(self, x: Float[Array, "dim"]):
         """Apply layer scaling to input tensor.
@@ -87,6 +86,54 @@ class LayerScale(eqx.Module):
         return x * self.gamma
 
 
+class DyT(eqx.Module):
+    """Dynamic Tanh layear.
+
+    This layer implements the DyT layer introduced in the Transformer
+    without Normalization paper[1].
+
+    Attributes:
+        init_values: Initial scale value (static)
+        gamma: Learnable scale parameters of size dim
+
+    References:
+        [1]. Zhu, et al., Transformers without Normalization. 2025.
+             https://arxiv.org/abs/2503.10622
+    """
+
+    alpha: Float[Array, "dim"]
+    weight: Float[Array, "dim"]
+    bias: Float[Array, "dim"]
+
+    def __init__(self, dim: int, alpha_init_value: float = 0.5):
+        """Initialize DyT.
+
+        Args:
+            dim: Dimension of the input features
+            alpha_init_value: Initial value for the scaling factor
+        """
+        self.alpha = jnp.repeat(alpha_init_value, dim)
+        self.weight = jnp.ones(dim)
+        self.bias = jnp.zeros(dim)
+
+    def __call__(
+        self,
+        x: Float[Array, "dim"],
+        *args,
+        **kwargs,
+    ):
+        """Apply dynamic tanh to input tensor.
+
+        Args:
+            x: Input tensor of shape (dim,)
+
+        Returns:
+            Scaled tensor of same shape as input
+        """
+        x = jnp.tanh(self.alpha * x)
+        return x * self.weight + self.bias
+
+
 def get_norm(module: str | eqx.Module) -> eqx.Module:
     """Get an `eqx.Module` from its common name.
 
@@ -107,5 +154,7 @@ def get_norm(module: str | eqx.Module) -> eqx.Module:
             return RMSNormGated
         case "layerscale":
             return LayerScale
+        case "dynamictanh":
+            return DyT
         case _:
             raise ValueError(f"Got an unknown module string: {module}")

{equimo-0.3.0 → equimo-0.3.2}/equimo/models/vit.py

@@ -403,6 +403,7 @@ class VisionTransformer(eqx.Module):
         key: PRNGKeyArray,
         mask: Optional[Int[Array, "embed_h embed_w"]] = None,
         inference: Optional[bool] = None,
+        **kwargs,
     ) -> Float[Array, "seqlen dim"]:
         """Extract features from input image.
 
@@ -434,7 +435,7 @@ class VisionTransformer(eqx.Module):
         x = self._pos_embed(x, h=self.embed_size, w=self.embed_size)
 
         for blk, key_block in zip(self.blocks, block_subkeys):
-            x = blk(x, inference=inference, key=key_block)
+            x = blk(x, inference=inference, key=key_block, **kwargs)
 
         return x
 
@@ -443,6 +444,7 @@ class VisionTransformer(eqx.Module):
         x: Float[Array, "channels height width"],
         key: PRNGKeyArray,
         inference: Optional[bool] = None,
+        **kwargs,
     ) -> dict:
         """Process features and return intermediate representations.
 
@@ -458,7 +460,7 @@ class VisionTransformer(eqx.Module):
                 - x_norm_patchtokens: Normalized patch tokens
                 - x_prenorm: Pre-normalized features
         """
-        x = self.features(x, inference=inference, key=key)
+        x = self.features(x, inference=inference, key=key, **kwargs)
         x_norm = jax.vmap(self.norm)(x)
 
         return {
@@ -473,6 +475,7 @@ class VisionTransformer(eqx.Module):
         x: Float[Array, "channels height width"],
         key: PRNGKeyArray = jr.PRNGKey(42),
         inference: Optional[bool] = None,
+        **kwargs,
     ) -> Float[Array, "num_classes"]:
         """Process input image through the full network.
 
@@ -484,7 +487,7 @@ class VisionTransformer(eqx.Module):
         Returns:
             Classification logits
         """
-        x = self.features(x, inference=inference, key=key)
+        x = self.features(x, inference=inference, key=key, **kwargs)
         x = jax.vmap(self.norm)(x)
         x = pool_sd(
             x,

{equimo-0.3.0 → equimo-0.3.2}/pyproject.toml

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

equimo-0.3.0/equimo/__init__.py

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