Equimo 0.4.0a15__tar.gz → 0.4.2__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: Equimo
-Version: 0.4.0a15
+Version: 0.4.2
 Summary: Implementation of popular vision models in Jax
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Developers
@@ -11,7 +11,7 @@ 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.10
+Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE.md
 Requires-Dist: einops>=0.8.0
@@ -78,6 +78,7 @@ Beyond normal ViT (e.g., dinov2 or siglip), equimo proposes other SotA architect
 | VSSD          | [VSSD: Vision Mamba with Non-Causal State Space Duality](https://arxiv.org/abs/2407.18559)                                                                      | 2024 | ✅        |
 | ReduceFormer  | [ReduceFormer: Attention with Tensor Reduction by Summation](https://arxiv.org/abs/2406.07488)                                                                  | 2024 | ✅        |
 | LowFormer     | [LowFormer: Hardware Efficient Design for Convolutional Transformer Backbones](https://arxiv.org/abs/2409.03460)                                                | 2024 | ✅        |
+| DINOv3        | [DINOv3](https://arxiv.org/abs/2508.10104)                                                                                                                      | 2025 | ✅        |
 
 \*: Only contains the [Linear Angular Attention](https://github.com/clementpoiret/Equimo/blob/f8fcc79e45ca65e9deb1d970c4286c0b8562f9c2/equimo/layers/attention.py#L1407) module. It is straight forward to build a ViT around it, but may require an additional `__call__` kwarg to control the `sparse_reg` bool.
 

{equimo-0.4.0a15 → equimo-0.4.2}/README.md

@@ -42,6 +42,7 @@ Beyond normal ViT (e.g., dinov2 or siglip), equimo proposes other SotA architect
 | VSSD          | [VSSD: Vision Mamba with Non-Causal State Space Duality](https://arxiv.org/abs/2407.18559)                                                                      | 2024 | ✅        |
 | ReduceFormer  | [ReduceFormer: Attention with Tensor Reduction by Summation](https://arxiv.org/abs/2406.07488)                                                                  | 2024 | ✅        |
 | LowFormer     | [LowFormer: Hardware Efficient Design for Convolutional Transformer Backbones](https://arxiv.org/abs/2409.03460)                                                | 2024 | ✅        |
+| DINOv3        | [DINOv3](https://arxiv.org/abs/2508.10104)                                                                                                                      | 2025 | ✅        |
 
 \*: Only contains the [Linear Angular Attention](https://github.com/clementpoiret/Equimo/blob/f8fcc79e45ca65e9deb1d970c4286c0b8562f9c2/equimo/layers/attention.py#L1407) module. It is straight forward to build a ViT around it, but may require an additional `__call__` kwarg to control the `sparse_reg` bool.
 

{equimo-0.4.0a15 → equimo-0.4.2}/pyproject.toml

@@ -1,9 +1,9 @@
 [project]
 name = "Equimo"
-version = "0.4.0-alpha.15"
+version = "0.4.2"
 description = "Implementation of popular vision models in Jax"
 readme = "README.md"
-requires-python = ">=3.10"
+requires-python = ">=3.11"
 dependencies = [
   "einops>=0.8.0",
   "equinox>=0.11.5",

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: Equimo
-Version: 0.4.0a15
+Version: 0.4.2
 Summary: Implementation of popular vision models in Jax
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Developers
@@ -11,7 +11,7 @@ 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.10
+Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE.md
 Requires-Dist: einops>=0.8.0
@@ -78,6 +78,7 @@ Beyond normal ViT (e.g., dinov2 or siglip), equimo proposes other SotA architect
 | VSSD          | [VSSD: Vision Mamba with Non-Causal State Space Duality](https://arxiv.org/abs/2407.18559)                                                                      | 2024 | ✅        |
 | ReduceFormer  | [ReduceFormer: Attention with Tensor Reduction by Summation](https://arxiv.org/abs/2406.07488)                                                                  | 2024 | ✅        |
 | LowFormer     | [LowFormer: Hardware Efficient Design for Convolutional Transformer Backbones](https://arxiv.org/abs/2409.03460)                                                | 2024 | ✅        |
+| DINOv3        | [DINOv3](https://arxiv.org/abs/2508.10104)                                                                                                                      | 2025 | ✅        |
 
 \*: Only contains the [Linear Angular Attention](https://github.com/clementpoiret/Equimo/blob/f8fcc79e45ca65e9deb1d970c4286c0b8562f9c2/equimo/layers/attention.py#L1407) module. It is straight forward to build a ViT around it, but may require an additional `__call__` kwarg to control the `sparse_reg` bool.
 

equimo-0.4.2/src/equimo/__init__.py

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

{equimo-0.4.0a15 → equimo-0.4.2}/src/equimo/conversion/utils.py

@@ -36,10 +36,11 @@ def convert_params_from_torch(
     replace_cfg: Dict[str, str],
     expand_cfg: Dict[str, list],
     squeeze_cfg: Dict[str, int | None],
-    whitelist: list[str],
+    torch_whitelist: list[str],
+    jax_whitelist: list[str],
     strict: bool = True,
     source: Literal["torchhub", "timm", "custom"] = "torchhub",
-    torch_hub_cfg: Optional[list[str]] = None,
+    torch_hub_cfg: Optional[dict] = None,
     torch_model=None,
     timm_cfg: Optional[list] = None,
     return_torch: bool = False,
@@ -49,14 +50,15 @@ def convert_params_from_torch(
 
     Args:
         jax_model (eqx.Module): A preexisting Jax model corresponding to the checkpoint to download.
-        torch_hub_cfg (Tuple[str]): Arguments passed to `torch.hub.load()`.
+        torch_hub_cfg (dict): Arguments passed to `torch.hub.load()`.
         replace_cfg (Dict[str, str]): Rename parameters from key to value.
         expand_cfg (Dict[str, list]): Config to reshape params, see `expand_torch_tensor`
         sqeeze_cfg (Dict[str, int|None]): Config to squeeze tensors, opposite of expand.
-        whitelist (Set[str]): Parameters to exclude from format conversion.
+        torch_whitelist (Set[str]): Parameters to exclude from format conversion.
+        jax_whitelist (Set[str]): Parameters to exclude from format conversion.
         strict (bool): Whether to crash on missing parameters one of the models.
         source (str): Torch Hub or timm.
-        torch_hub_cfg (Optional[list]): args to pass to `torch.hub.load`.
+        torch_hub_cfg (dict): args to pass to `torch.hub.load`.
         torch_model [torch.nn.Module]: Custom torch model
         timm_cfg (Optional[list]): args to pass to `timm.create_model`.
         return_torch (bool): Return both jax and torch models.
@@ -79,7 +81,7 @@ def convert_params_from_torch(
                 raise ValueError(
                     "The `torchhub` source is selected but `torch_hub_cfg` is None."
                 )
-            torch_model = torch.hub.load(*torch_hub_cfg)
+            torch_model = torch.hub.load(**torch_hub_cfg)
         case "timm":
             if timm_cfg is None:
                 raise ValueError(
@@ -107,12 +109,20 @@ def convert_params_from_torch(
 
         if param_path not in torch_params:
             _msg = f"{param_path} ({shape}) not found in PyTorch model."
-            if strict:
+            if strict and param_path not in jax_whitelist:
                 logger.error(_msg)
                 raise AttributeError(_msg)
 
-            logger.warning(f"{_msg} Appending `None` to flat param list.")
-            torch_params_flat.append(None)
+            if param_path in jax_whitelist:
+                p = param
+                logger.warning(
+                    f"{_msg} Appending original parameters to flat param list because of `jax_whitelist`."
+                )
+            else:
+                p = None
+                logger.warning(f"{_msg} Appending `None` to flat param list.")
+
+            torch_params_flat.append(p)
             continue
 
         logger.info(f"Converting {param_path}...")
@@ -137,7 +147,7 @@ def convert_params_from_torch(
         logger.warning(
             f"PyTorch parameters `{path}` ({param.shape}) were not converted."
         )
-        if strict and path not in whitelist:
+        if strict and path not in torch_whitelist:
             _msg = f"The PyTorch model contains parameters ({path}) that do not have a Jax counterpart."
             logger.error(_msg)
             raise AttributeError(_msg)
@@ -152,10 +162,11 @@ def convert_torch_to_equinox(
     replace_cfg: dict = {},
     expand_cfg: dict = {},
     squeeze_cfg: dict = {},
-    whitelist: list[str] = [],
+    torch_whitelist: list[str] = [],
+    jax_whitelist: list[str] = [],
     strict: bool = True,
     source: Literal["torchhub", "timm"] = "torchhub",
-    torch_hub_cfg: Optional[list[str]] = None,
+    torch_hub_cfg: Optional[dict] = None,
     torch_model=None,
     timm_cfg: Optional[list] = None,
     return_torch: bool = False,
@@ -168,10 +179,11 @@ def convert_torch_to_equinox(
         replace_cfg: Dict of parameter name replacements
         expand_cfg: Dict of dimensions to expand
         squeeze_cfg: Dict of dimensions to squeeze
-        whitelist: List of parameters to keep from JAX model
+        torch_whitelist: List of parameters allowed to be in PT model but not in Jax
+        jax_whitelist: List of parameters allowed to be in Jax model but not in PT
         strict: Wether to raise an issue if not all weights are converted
         source (str): Torch Hub or timm.
-        torch_hub_cfg: [repo, model_name] for torch.hub.load
+        torch_hub_cfg (dict): torch.hub.load config
         torch_model [torch.nn.Module]: Custom torch model
         timm_cfg (Optional[list]): args to pass to `timm.create_model`.
         return_torch (bool): Return both jax and torch models.
@@ -186,7 +198,8 @@ def convert_torch_to_equinox(
             replace_cfg,
             expand_cfg,
             squeeze_cfg,
-            whitelist,
+            torch_whitelist,
+            jax_whitelist,
             strict,
             source,
             torch_hub_cfg,
@@ -204,7 +217,8 @@ def convert_torch_to_equinox(
             replace_cfg,
             expand_cfg,
             squeeze_cfg,
-            whitelist,
+            torch_whitelist,
+            jax_whitelist,
             strict,
             source,
             torch_hub_cfg,

{equimo-0.4.0a15 → equimo-0.4.2}/src/equimo/io.py

@@ -25,7 +25,7 @@ def save_model(
     path: Path,
     model: eqx.Module,
     model_config: dict,
-    torch_hub_cfg: list[str] = [],
+    torch_hub_cfg: list[str] | dict = {},
     timm_cfg: list = [],
     compression: bool = True,
 ):
@@ -245,11 +245,32 @@ def load_model(
     return model
 
 
+def _center_crop_square(array: jnp.ndarray) -> jnp.ndarray:
+    """Center-crop a HxW(xC) array to a square of side min(H, W).
+
+    Args:
+        array: jnp.ndarray with shape (H, W) or (H, W, C).
+
+    Returns:
+        Center-cropped array with shape (M, M) or (M, M, C) where M = min(H, W).
+    """
+    if array.ndim < 2:
+        raise ValueError("Input array must have at least 2 dimensions (H, W[, C]).")
+    h, w = array.shape[:2]
+    if h == w:
+        return array
+    m = min(h, w)
+    top = (h - m) // 2
+    left = (w - m) // 2
+    return array[top : top + m, left : left + m, ...]
+
+
 def load_image(
     path: str,
     mean: Optional[list[float]] = None,
     std: Optional[list[float]] = None,
     size: Optional[int] = None,
+    center_crop: bool = False,
 ):
     """Load an image and perform minor preprocessing.
 
@@ -258,6 +279,7 @@ def load_image(
         mean (list, optional): Channel mean for normalization.
         std (list, optional): Channel std for normalization.
         size (int, optional): Size to which resize the image.
+        center_crop (bool, optional): If True, center-crop to square prior to resizing.
 
     Returns:
         jnp.array: The loaded image.
@@ -275,6 +297,10 @@ def load_image(
         pil_image = Image.open(image_bytes)
 
         array = jnp.array(pil_image).astype(jnp.float32) / 255.0
+
+        if center_crop:
+            array = _center_crop_square(array)
+
         if size is not None:
             array = jax.image.resize(array, (size, size, 3), method="bilinear")
 

{equimo-0.4.0a15 → equimo-0.4.2}/src/equimo/layers/attention.py

@@ -16,6 +16,70 @@ from equimo.layers.posemb import PosCNN2D, PosEmbMLPSwinv1D, PosEmbMLPSwinv2D, R
 from equimo.utils import nearest_power_of_2_divisor
 
 
+def rope_rotate_half(x: jax.Array) -> jax.Array:
+    """Rotate last-dim pairs by 90 degrees: [x0..x_{D/2-1}, x_{D/2}..x_{D-1}]
+    -> [-x_{D/2}..-x_{D-1}, x0..x_{D/2-1}]
+    """
+    x1, x2 = jnp.split(x, 2, axis=-1)
+    return jnp.concatenate([-x2, x1], axis=-1)
+
+
+def rope_apply(x: jax.Array, sin: jax.Array, cos: jax.Array) -> jax.Array:
+    """Apply RoPE to `x` using per-position sin/cos.
+
+    Shapes:
+    - x:   [..., D]
+    - sin: [..., D]
+    - cos: [..., D]
+    Broadcasting across leading axes is supported.
+    """
+    return (x * cos) + (rope_rotate_half(x) * sin)
+
+
+def rope_apply_qk_last_hw(
+    q: jax.Array, k: jax.Array, sin: jax.Array, cos: jax.Array
+) -> tuple[jax.Array, jax.Array]:
+    """Apply RoPE to the last HW tokens of q and k (keeping any prefix unchanged).
+
+    Shapes:
+    - q, k: [H, N, D] (heads, tokens, head_dim)
+    - sin, cos: [HW, D]
+    - N = prefix + HW
+
+    Returns:
+    - q_rot, k_rot: same shape as inputs.
+    """
+    if sin.shape[-1] != q.shape[-1] or cos.shape[-1] != q.shape[-1]:
+        raise ValueError(
+            f"sin/cos last dim must equal head_dim; got {sin.shape[-1]} vs {q.shape[-1]}"
+        )
+    N = q.shape[-2]
+    HW = sin.shape[-2]
+    prefix = N - HW
+    if prefix < 0:
+        raise ValueError(f"Sequence length N={N} smaller than HW={HW}.")
+
+    q_dtype, k_dtype = q.dtype, k.dtype
+    rope_dtype = sin.dtype
+
+    q = q.astype(rope_dtype)
+    k = k.astype(rope_dtype)
+
+    # Broadcast sin/cos across heads
+    sin_b = sin[None, :, :]  # [1, HW, D]
+    cos_b = cos[None, :, :]  # [1, HW, D]
+
+    q_prefix, q_tail = jnp.split(q, [prefix], axis=-2)  # axis -2 is tokens
+    k_prefix, k_tail = jnp.split(k, [prefix], axis=-2)
+
+    q_tail = rope_apply(q_tail, sin_b, cos_b)
+    k_tail = rope_apply(k_tail, sin_b, cos_b)
+
+    q_out = jnp.concatenate([q_prefix, q_tail], axis=-2).astype(q_dtype)
+    k_out = jnp.concatenate([k_prefix, k_tail], axis=-2).astype(k_dtype)
+    return q_out, k_out
+
+
 class Attention(eqx.Module):
     """Multi-head self attention module.
 
@@ -76,6 +140,7 @@ class Attention(eqx.Module):
         key: PRNGKeyArray,
         inference: Optional[bool] = None,
         mask: Optional[Float[Array, ""]] = None,
+        rope_sincos: Optional[Tuple[jax.Array, jax.Array]] = None,
     ) -> Float[Array, "seqlen dim"]:
         key1, key2 = jr.split(key, 2)
 
@@ -91,6 +156,15 @@ class Attention(eqx.Module):
         q = jax.vmap(jax.vmap(self.q_norm))(q)
         k = jax.vmap(jax.vmap(self.k_norm))(k)
 
+        if rope_sincos is not None:
+            sin, cos = rope_sincos  # [HW, D], [HW, D]
+            if sin.shape[-1] != self.head_dim or cos.shape[-1] != self.head_dim:
+                raise ValueError(
+                    f"RoPE sin/cos last dim ({sin.shape[-1]}) must equal head_dim ({self.head_dim})."
+                )
+            # leave any prefix tokens untouched; rotate last HW tokens
+            q, k = rope_apply_qk_last_hw(q, k, sin, cos)
+
         attn = jnp.einsum("hqd,hkd->hqk", q, k) / jnp.sqrt(self.head_dim)
 
         if mask is not None:
@@ -255,6 +329,7 @@ class AttentionBlock(eqx.Module):
         ffn_layer: eqx.Module = Mlp,
         ffn_bias: bool = True,
         ffn_norm: bool = False,
+        ffn_kwargs: dict = {},
         norm_layer: eqx.Module = eqx.nn.LayerNorm,
         post_attention_norm: bool = False,
         init_values: float | None = None,
@@ -308,6 +383,7 @@ class AttentionBlock(eqx.Module):
             bias=ffn_bias,
             eps=eps,
             key=key_mlp,
+            **ffn_kwargs,
         )
 
         self.drop_path1 = DropPathAdd(dr1)
@@ -317,6 +393,7 @@ class AttentionBlock(eqx.Module):
         self,
         x: Float[Array, "seqlen dim"],
         key: PRNGKeyArray,
+        rope_sincos: Optional[Tuple[jax.Array, jax.Array]] = None,
         inference: Optional[bool] = None,
         mask: Optional[Float[Array, ""]] = None,
     ) -> Float[Array, "seqlen dim"]:
@@ -325,6 +402,11 @@ class AttentionBlock(eqx.Module):
         # I chose to define extra args here rather than passing mask directly
         # because not all attention mechanisms support masks as args
         extra_kwargs = {"mask": mask} if mask is not None else {}
+        attn_kwargs = (
+            extra_kwargs | {"rope_sincos": rope_sincos}
+            if rope_sincos is not None
+            else extra_kwargs
+        )
 
         x = self.drop_path1(
             x,
@@ -334,7 +416,7 @@ class AttentionBlock(eqx.Module):
                         jax.vmap(self.prenorm)(x),
                         inference=inference,
                         key=key_attn,
-                        **extra_kwargs,
+                        **attn_kwargs,
                     )
                 )
             ),

{equimo-0.4.0a15 → equimo-0.4.2}/src/equimo/layers/convolution.py

@@ -900,10 +900,10 @@ class DSConv(eqx.Module):
         key: PRNGKeyArray,
         inference: Optional[bool] = None,
     ):
-        key_dropout, key_droppath = jr.split(key, 2)
+        key_depth, key_point, key_dropout, key_droppath = jr.split(key, 4)
 
-        out = self.depth_conv(x)
-        out = self.point_conv(out)
+        out = self.depth_conv(x, inference=inference, key=key_depth)
+        out = self.point_conv(out, inference=inference, key=key_point)
 
         out = self.dropout(out, inference=inference, key=key_dropout)
 

{equimo-0.4.0a15 → equimo-0.4.2}/src/equimo/layers/ffn.py

@@ -244,14 +244,8 @@ class SwiGlu(eqx.Module):
     a gating mechanism where the input is transformed by two parallel paths and
     combined multiplicatively.
 
-    The computation flow is:
-    1. Joint projection to higher dimension (w12)
-    2. Split into two paths
-    3. Apply SiLU to first path and multiply with second path
-    4. Project back to original dimension (w3)
-
     Attributes:
-        w12: Joint projection layer for both paths
+        w1, w2: projection layers for both paths
         w3: Final projection layer
         drop1: Dropout after gating
         drop2: Dropout after final projection
@@ -260,7 +254,8 @@ class SwiGlu(eqx.Module):
         [1]: https://arxiv.org/pdf/2002.05202
     """
 
-    w12: eqx.nn.Linear
+    w1: eqx.nn.Linear
+    w2: eqx.nn.Linear
     w3: eqx.nn.Linear
     drop1: eqx.nn.Dropout
     drop2: eqx.nn.Dropout
@@ -273,6 +268,7 @@ class SwiGlu(eqx.Module):
         out_features: int | None = None,
         hidden_features: int | None = None,
         dropout_rate: float = 0.0,
+        align_to: int = 8,
         bias: bool = True,
         **kwargs,
     ):
@@ -284,19 +280,25 @@ class SwiGlu(eqx.Module):
             out_features: Number of output features (default: same as in_features)
             hidden_features: Size of hidden dimension (default: same as in_features)
             dropout_rate: Dropout probability (default: 0.0)
+            align_to: constrains hidden features to be a multiple of a given int (default: 8)
             bias: Whether to include bias in linear layers (default: True)
             **kwargs: Additional arguments
         """
-        key_fc1, key_fc2 = jr.split(key, 2)
+        key_fc1, key_fc2, key_fc3 = jr.split(key, 3)
 
-        hidden_features = hidden_features or in_features
         out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        d = int(hidden_features * 2 / 3)
+        hidden_features = d + (-d % align_to)
 
-        self.w12 = eqx.nn.Linear(
-            in_features, 2 * hidden_features, use_bias=bias, key=key_fc1
+        self.w1 = eqx.nn.Linear(
+            in_features, hidden_features, use_bias=bias, key=key_fc1
+        )
+        self.w2 = eqx.nn.Linear(
+            in_features, hidden_features, use_bias=bias, key=key_fc2
         )
         self.w3 = eqx.nn.Linear(
-            hidden_features, out_features, use_bias=bias, key=key_fc2
+            hidden_features, out_features, use_bias=bias, key=key_fc3
         )
 
         self.drop1 = eqx.nn.Dropout(dropout_rate)
@@ -310,8 +312,9 @@ class SwiGlu(eqx.Module):
     ) -> Float[Array, "seqlen dim"]:
         key_dr1, key_dr2 = jr.split(key, 2)
 
-        x12 = jax.vmap(self.w12)(x)
-        x1, x2 = jnp.split(x12, 2, axis=-1)
+        x1 = jax.vmap(self.w1)(x)
+        x2 = jax.vmap(self.w2)(x)
+
         x = self.drop1(
             jax.nn.silu(x1) * x2,
             inference=inference,
@@ -327,7 +330,38 @@ class SwiGlu(eqx.Module):
         return x
 
 
-class SwiGluFused(SwiGlu):
+class SwiGluFused(eqx.Module):
+    """SwiGLU activation module with dropout.
+
+    This matches the implementation of Dinov2 giant at
+    https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/layers/swiglu_ffn.py#L54
+
+    Implements the SwiGLU (Swish-Gated Linear Unit) activation function with dropout,
+    as described in "GLU Variants Improve Transformer" paper [1]. The architecture uses
+    a gating mechanism where the input is transformed by two parallel paths and
+    combined multiplicatively.
+
+    The computation flow is:
+    1. Joint projection to higher dimension (w12)
+    2. Split into two paths
+    3. Apply SiLU to first path and multiply with second path
+    4. Project back to original dimension (w3)
+
+    Attributes:
+        w12: Joint projection layer for both paths
+        w3: Final projection layer
+        drop1: Dropout after gating
+        drop2: Dropout after final projection
+
+    References:
+        [1]: https://arxiv.org/pdf/2002.05202
+    """
+
+    w12: eqx.nn.Linear
+    w3: eqx.nn.Linear
+    drop1: eqx.nn.Dropout
+    drop2: eqx.nn.Dropout
+
     def __init__(
         self,
         in_features: int,
@@ -335,27 +369,61 @@ class SwiGluFused(SwiGlu):
         key: PRNGKeyArray,
         out_features: int | None = None,
         hidden_features: int | None = None,
-        dropout_rate: float = 0,
+        dropout_rate: float = 0.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
+        """Initialize the SwiGLU module.
+
+        Args:
+            in_features: Number of input features
+            key: PRNG key for initialization
+            out_features: Number of output features (default: same as in_features)
+            hidden_features: Size of hidden dimension (default: same as in_features)
+            dropout_rate: Dropout probability (default: 0.0)
+            bias: Whether to include bias in linear layers (default: True)
+            **kwargs: Additional arguments
         """
+        key_fc1, key_fc2 = jr.split(key, 2)
+
         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,
+        self.w12 = eqx.nn.Linear(
+            in_features, 2 * hidden_features, use_bias=bias, key=key_fc1
+        )
+        self.w3 = eqx.nn.Linear(
+            hidden_features, out_features, use_bias=bias, key=key_fc2
+        )
+
+        self.drop1 = eqx.nn.Dropout(dropout_rate)
+        self.drop2 = eqx.nn.Dropout(dropout_rate)
+
+    def __call__(
+        self,
+        x: Float[Array, "seqlen dim"],
+        key: PRNGKeyArray,
+        inference: Optional[bool] = None,
+    ) -> Float[Array, "seqlen dim"]:
+        key_dr1, key_dr2 = jr.split(key, 2)
+
+        x12 = jax.vmap(self.w12)(x)
+        x1, x2 = jnp.split(x12, 2, axis=-1)
+        x = self.drop1(
+            jax.nn.silu(x1) * x2,
+            inference=inference,
+            key=key_dr1,
         )
 
+        x = self.drop2(
+            jax.vmap(self.w3)(x),
+            inference=inference,
+            key=key_dr2,
+        )
+
+        return x
+
 
 def get_ffn(module: str | eqx.Module) -> eqx.Module:
     """Get an `eqx.Module` from its common name.