Equimo 0.3.2__tar.gz → 0.3.3__tar.gz

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

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

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

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

equimo-0.3.3/equimo/__init__.py

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

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

@@ -19,7 +19,11 @@ class RMSNormGated(eqx.Module):
 
     w: Float[Array, "dim"]
 
-    def __init__(self, d: int):
+    def __init__(
+        self,
+        d: int,
+        **kwargs,
+    ):
         """Initialize RMSNormGated.
 
         Args:
@@ -31,8 +35,6 @@ 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.
 
@@ -105,7 +107,12 @@ class DyT(eqx.Module):
     weight: Float[Array, "dim"]
     bias: Float[Array, "dim"]
 
-    def __init__(self, dim: int, alpha_init_value: float = 0.5):
+    def __init__(
+        self,
+        dim: int,
+        alpha_init_value: float = 0.5,
+        **kwargs,
+    ):
         """Initialize DyT.
 
         Args:
@@ -119,8 +126,6 @@ class DyT(eqx.Module):
     def __call__(
         self,
         x: Float[Array, "dim"],
-        *args,
-        **kwargs,
     ):
         """Apply dynamic tanh to input tensor.
 

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

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

equimo-0.3.2/equimo/__init__.py

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