Equimo 0.2.0__tar.gz → 0.2.1__tar.gz

{equimo-0.2.0 → equimo-0.2.1}/Equimo.egg-info/PKG-INFO

@@ -1,7 +1,16 @@
 Metadata-Version: 2.2
 Name: Equimo
-Version: 0.2.0
+Version: 0.2.1
 Summary: Implementation of popular vision models in Jax
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 Requires-Dist: einops>=0.8.0

{equimo-0.2.0 → equimo-0.2.1}/PKG-INFO

@@ -1,7 +1,16 @@
 Metadata-Version: 2.2
 Name: Equimo
-Version: 0.2.0
+Version: 0.2.1
 Summary: Implementation of popular vision models in Jax
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 Requires-Dist: einops>=0.8.0

{equimo-0.2.0 → equimo-0.2.1}/equimo/layers/attention.py

@@ -1552,3 +1552,53 @@ class LinearAngularAttention(eqx.Module):
         x = self.proj_drop(x, inference=inference, key=key2)
 
         return x
+
+
+def get_attention(module: str | eqx.Module) -> eqx.Module:
+    """Get an `eqx.Module` from its common name.
+
+    This is necessary because configs have to be stringified and stored as
+    json files to allow (de)serialization.
+    """
+    if not isinstance(module, str):
+        return module
+
+    match module:
+        case "attention":
+            return Attention
+        case "windowedattention":
+            return WindowedAttention
+        case "shsa":
+            return SHSA
+        case "linearattention":
+            return LinearAttention
+        case "mmsa":
+            return MMSA
+        case "sqa":
+            return SQA
+        case "linearangularattention":
+            return LinearAngularAttention
+        case _:
+            raise ValueError(f"Got an unknown module string: {module}")
+
+
+def get_attention_block(module: str | eqx.Module) -> eqx.Module:
+    """Get an `eqx.Module` from its common name.
+
+    This is necessary because configs have to be stringified and stored as
+    json files to allow (de)serialization.
+    """
+    if not isinstance(module, str):
+        return module
+
+    match module:
+        case "attentionblock":
+            return AttentionBlock
+        case "hatblock":
+            return HATBlock
+        case "mllablock":
+            return MllaBlock
+        case "partialformerblock":
+            return PartialFormerBlock
+        case _:
+            raise ValueError(f"Got an unknown module string: {module}")

{equimo-0.2.0 → equimo-0.2.1}/equimo/layers/ffn.py

@@ -285,10 +285,10 @@ class SwiGlu(eqx.Module):
         out_features = out_features or in_features
 
         self.w12 = eqx.nn.Linear(
-            in_features, hidden_features, use_bias=bias, key=key_fc1
+            in_features, 2 * hidden_features, use_bias=bias, key=key_fc1
         )
         self.w3 = eqx.nn.Linear(
-            hidden_features // 2, out_features, use_bias=bias, key=key_fc2
+            hidden_features, out_features, use_bias=bias, key=key_fc2
         )
 
         self.drop1 = eqx.nn.Dropout(dropout_rate)
@@ -317,3 +317,57 @@ class SwiGlu(eqx.Module):
         )
 
         return x
+
+
+class SwiGluFused(SwiGlu):
+    def __init__(
+        self,
+        in_features: int,
+        *,
+        key: PRNGKeyArray,
+        out_features: int | None = None,
+        hidden_features: int | None = None,
+        dropout_rate: float = 0,
+        bias: bool = True,
+        **kwargs,
+    ):
+        """This matches the implementation of Dinov2 giant at
+        https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/layers/swiglu_ffn.py#L54
+        """
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8
+
+        super().__init__(
+            in_features,
+            key=key,
+            out_features=out_features,
+            hidden_features=hidden_features,
+            dropout_rate=dropout_rate,
+            bias=bias,
+            **kwargs,
+        )
+
+
+def get_ffn(module: str | eqx.Module) -> eqx.Module:
+    """Get an `eqx.Module` from its common name.
+
+    This is necessary because configs have to be stringified and stored as
+    json files to allow (de)serialization.
+    """
+    if not isinstance(module, str):
+        return module
+
+    match module:
+        case "mlp":
+            return Mlp
+        case "swiglu":
+            return SwiGlu
+        case "swiglufused":
+            return SwiGluFused
+        case "dinohead":
+            return DINOHead
+        case "weightnormlinear":
+            return WeightNormLinear
+        case _:
+            raise ValueError(f"Got an unknown module string: {module}")

{equimo-0.2.0 → equimo-0.2.1}/equimo/layers/norm.py

@@ -85,3 +85,27 @@ class LayerScale(eqx.Module):
             Scaled tensor of same shape as input
         """
         return x * self.gamma
+
+
+def get_norm(module: str | eqx.Module) -> eqx.Module:
+    """Get an `eqx.Module` from its common name.
+
+    This is necessary because configs have to be stringified and stored as
+    json files to allow (de)serialization.
+    """
+    if not isinstance(module, str):
+        return module
+
+    match module:
+        case "layernorm":
+            return eqx.nn.LayerNorm
+        case "rmsnorm":
+            return eqx.nn.RMSNorm
+        case "groupnorm":
+            return eqx.nn.GroupNorm
+        case "rmsnormgated":
+            return RMSNormGated
+        case "layerscale":
+            return LayerScale
+        case _:
+            raise ValueError(f"Got an unknown module string: {module}")

{equimo-0.2.0 → equimo-0.2.1}/equimo/models/vit.py

@@ -8,8 +8,14 @@ import jax.random as jr
 from einops import rearrange
 from jaxtyping import Array, Float, Int, PRNGKeyArray
 
-from equimo.layers.attention import Attention, AttentionBlock
-from equimo.layers.ffn import Mlp
+from equimo.layers.attention import (
+    Attention,
+    AttentionBlock,
+    get_attention,
+    get_attention_block,
+)
+from equimo.layers.ffn import Mlp, get_ffn
+from equimo.layers.norm import get_norm
 from equimo.layers.patch import PatchEmbedding
 from equimo.layers.posemb import PosCNN
 from equimo.utils import pool_sd, to_list
@@ -177,7 +183,7 @@ class VisionTransformer(eqx.Module):
         pos_drop_rate: float = 0.0,
         drop_path_rate: float = 0.0,
         drop_path_uniform: bool = False,
-        block: eqx.Module = AttentionBlock,
+        block: str | eqx.Module = AttentionBlock,
         mlp_ratio: float = 4.0,
         qkv_bias: bool = True,
         proj_bias: bool = True,
@@ -185,10 +191,10 @@ class VisionTransformer(eqx.Module):
         attn_drop: float = 0.0,
         proj_drop: float = 0.0,
         act_layer: Callable = jax.nn.gelu,
-        attn_layer: eqx.Module = Attention,
-        ffn_layer: eqx.Module = Mlp,
+        attn_layer: str | eqx.Module = Attention,
+        ffn_layer: str | eqx.Module = Mlp,
         ffn_bias: bool = True,
-        norm_layer: eqx.Module = eqx.nn.LayerNorm,
+        norm_layer: str | eqx.Module = eqx.nn.LayerNorm,
         init_values: float | None = None,
         global_pool: Literal["", "token", "avg", "avgmax", "max"] = "avg",
         num_classes: int = 1000,
@@ -212,6 +218,11 @@ class VisionTransformer(eqx.Module):
         self.global_pool = global_pool
         self.embed_size = img_size // patch_size
 
+        block = get_attention_block(block)
+        attn_layer = get_attention(attn_layer)
+        ffn_layer = get_ffn(ffn_layer)
+        norm_layer = get_norm(norm_layer)
+
         self.patch_embed = PatchEmbedding(
             in_channels,
             dim,

{equimo-0.2.0 → equimo-0.2.1}/equimo/utils.py

@@ -1,4 +1,5 @@
 from functools import partial
+
 import jax
 import jax.numpy as jnp
 from jaxtyping import Array, Float

{equimo-0.2.0 → equimo-0.2.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "Equimo"
-version = "0.2.0"
+version = "0.2.1"
 description = "Implementation of popular vision models in Jax"
 readme = "README.md"
 requires-python = ">=3.11"
@@ -13,6 +13,21 @@ dependencies = [
   "lz4>=4.4.3",
   "requests>=2.32.3",
 ]
+classifiers = [
+  "Development Status :: 4 - Beta",
+
+  "Intended Audience :: Developers",
+  "Intended Audience :: Science/Research",
+
+  "License :: OSI Approved :: MIT License",
+
+  "Topic :: Scientific/Engineering :: Artificial Intelligence",
+
+  "Programming Language :: Python :: 3.11",
+  "Programming Language :: Python :: 3.12",
+  "Programming Language :: Python :: 3.13",
+  "Programming Language :: Python :: 3.14",
+]
 
 [project.optional-dependencies]
 torch = [