Equimo 0.4.6a4__tar.gz → 0.4.7__tar.gz

{equimo-0.4.6a4/src/Equimo.egg-info → equimo-0.4.7}/PKG-INFO

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

{equimo-0.4.6a4 → equimo-0.4.7}/pyproject.toml

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

{equimo-0.4.6a4 → equimo-0.4.7/src/Equimo.egg-info}/PKG-INFO

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

equimo-0.4.7/src/equimo/__init__.py

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

{equimo-0.4.6a4 → equimo-0.4.7}/src/equimo/experimental/text.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import Callable, Optional, Sequence
+from typing import Callable, Optional, Sequence, Tuple
 from urllib.parse import urlparse
 
 import equinox as eqx
@@ -136,7 +136,7 @@ class Transformer(eqx.Module):
         blocks: A list of `AttentionBlock` instances forming the transformer stack.
     """
 
-    blocks: list[AttentionBlock]
+    blocks: Tuple[AttentionBlock, ...]
 
     def __init__(
         self,
@@ -164,7 +164,7 @@ class Transformer(eqx.Module):
 
         act_layer = get_act(act_layer)
 
-        self.blocks = [
+        self.blocks = tuple(
             AttentionBlock(
                 dim=dim,
                 num_heads=num_heads,
@@ -173,7 +173,7 @@ class Transformer(eqx.Module):
                 key=keys[i],
             )
             for i in range(depth)
-        ]
+        )
 
     def __call__(
         self,

{equimo-0.4.6a4 → equimo-0.4.7}/src/equimo/layers/attention.py

@@ -1671,7 +1671,7 @@ class RFAttention(eqx.Module):
     eps: float = eqx.field(static=True)
 
     qkv: eqx.nn.Conv2d
-    aggreg: list[eqx.nn.Conv2d]
+    aggreg: Tuple[eqx.nn.Conv2d, ...]
     proj: SingleConvBlock
 
     def __init__(
@@ -1709,7 +1709,7 @@ class RFAttention(eqx.Module):
             use_bias=use_bias,
             key=key_qkv,
         )
-        self.aggreg = [
+        self.aggreg = tuple(
             eqx.nn.Conv2d(
                 in_channels=3 * total_dim,
                 out_channels=3 * total_dim,
@@ -1720,7 +1720,7 @@ class RFAttention(eqx.Module):
                 use_bias=use_bias,
             )
             for scale in scales
-        ]
+        )
         # TODO: test different normalizations
         self.proj = SingleConvBlock(
             in_channels=self.total_dim,

{equimo-0.4.6a4 → equimo-0.4.7}/src/equimo/layers/convolution.py

@@ -409,7 +409,7 @@ class C2f(eqx.Module):
 
     conv1: SingleConvBlock
     conv2: SingleConvBlock
-    blocks: list[ConvBottleneck]
+    blocks: Tuple[ConvBottleneck, ...]
 
     def __init__(
         self,
@@ -445,7 +445,7 @@ class C2f(eqx.Module):
             key=key_conv2,
         )
 
-        self.blocks = [
+        self.blocks = tuple(
             ConvBottleneck(
                 in_channels=self.hidden_channels,
                 out_channels=self.hidden_channels,
@@ -456,7 +456,7 @@ class C2f(eqx.Module):
                 key=key_blocks[i],
             )
             for i in range(n)
-        ]
+        )
 
     def __call__(
         self,
@@ -589,7 +589,7 @@ class C3k2(eqx.Module):
 
     conv1: SingleConvBlock
     conv2: SingleConvBlock
-    blocks: list[ConvBottleneck] | list[C3k]
+    blocks: Tuple[ConvBottleneck, ...] | Tuple[C3k, ...]
 
     def __init__(
         self,
@@ -627,7 +627,7 @@ class C3k2(eqx.Module):
         )
 
         if c3k:
-            self.blocks = [
+            self.blocks = tuple(
                 C3k(
                     in_channels=self.hidden_channels,
                     out_channels=self.hidden_channels,
@@ -637,9 +637,9 @@ class C3k2(eqx.Module):
                     key=key_blocks[i],
                 )
                 for i in range(n)
-            ]
+            )
         else:
-            self.blocks = [
+            self.blocks = tuple(
                 ConvBottleneck(
                     in_channels=self.hidden_channels,
                     out_channels=self.hidden_channels,
@@ -648,7 +648,7 @@ class C3k2(eqx.Module):
                     key=key_blocks[i],
                 )
                 for i in range(n)
-            ]
+            )
 
     def __call__(
         self,
@@ -1078,12 +1078,12 @@ class GenericGhostModule(eqx.Module):
     cheap_operation: eqx.nn.Conv2d
 
     # Training
-    primary_rpr_conv: list[eqx.nn.Conv2d]
+    primary_rpr_conv: Tuple[eqx.nn.Conv2d, ...]
     primary_rpr_scale: eqx.nn.Conv2d | eqx.nn.Identity
     primary_shared_norm: eqx.nn.GroupNorm
     primary_activation: Callable
 
-    cheap_rpr_conv: list[eqx.nn.Conv2d]
+    cheap_rpr_conv: Tuple[eqx.nn.Conv2d, ...]
     cheap_rpr_scale: eqx.nn.Conv2d | eqx.nn.Identity
     cheap_shared_norm: eqx.nn.GroupNorm
     cheap_activation: Callable
@@ -1156,7 +1156,7 @@ class GenericGhostModule(eqx.Module):
 
         # Primary training branches
         init_num_groups = nearest_power_of_2_divisor(init_channels, 32)
-        self.primary_rpr_conv = [
+        self.primary_rpr_conv = tuple(
             eqx.nn.Conv2d(
                 in_channels=in_channels,
                 out_channels=init_channels,
@@ -1167,7 +1167,7 @@ class GenericGhostModule(eqx.Module):
                 key=key_ps[i],
             )
             for i in range(num_conv_branches)
-        ]
+        )
         self.primary_rpr_scale = (
             eqx.nn.Conv2d(
                 in_channels=in_channels,
@@ -1186,7 +1186,7 @@ class GenericGhostModule(eqx.Module):
 
         # Cheap training branches (depthwise)
         newchannels_num_groups = nearest_power_of_2_divisor(new_channels, 32)
-        self.cheap_rpr_conv = [
+        self.cheap_rpr_conv = tuple(
             eqx.nn.Conv2d(
                 in_channels=init_channels,
                 out_channels=new_channels,
@@ -1198,7 +1198,7 @@ class GenericGhostModule(eqx.Module):
                 key=key_cs[i],
             )
             for i in range(self.num_conv_branches)
-        ]
+        )
         self.cheap_rpr_scale = (
             eqx.nn.Conv2d(
                 in_channels=init_channels,
@@ -1344,7 +1344,7 @@ class GhostBottleneck(eqx.Module):
     ghost2: "GenericGhostModule"
 
     dw_conv: eqx.nn.Conv2d | eqx.nn.Identity
-    dw_rpr_conv: list[eqx.nn.Conv2d]  # depthwise conv branches (no bias)
+    dw_rpr_conv: Tuple[eqx.nn.Conv2d, ...]  # depthwise conv branches (no bias)
     dw_rpr_scale: eqx.nn.Conv2d | eqx.nn.Identity  # optional 1x1 depthwise (no bias)
     dw_shared_norm: eqx.nn.GroupNorm | eqx.nn.Identity
 
@@ -1393,7 +1393,7 @@ class GhostBottleneck(eqx.Module):
         # Depthwise stage (only if stride > 1)
         if stride > 1:
             # Training-time branches (depthwise, no bias); no activation; shared GN after sum
-            self.dw_rpr_conv = [
+            self.dw_rpr_conv = tuple(
                 eqx.nn.Conv2d(
                     in_channels=mid_channels,
                     out_channels=mid_channels,
@@ -1405,7 +1405,7 @@ class GhostBottleneck(eqx.Module):
                     key=k_dw_list[i],
                 )
                 for i in range(3)
-            ]
+            )
             # Optional scale branch (1x1, depthwise, stride=stride)
             self.dw_rpr_scale = (
                 eqx.nn.Conv2d(
@@ -1543,15 +1543,6 @@ class GhostBottleneck(eqx.Module):
         return x
 
 
-key = jr.PRNGKey(42)
-x = jr.normal(key, (64, 16, 16))
-f = GhostBottleneck(64, 128, 64, se_ratio=0.25, key=key)
-x1 = f.ghost1(x, key=key)
-x2 = f.se(x1)
-x3 = f.ghost2(x2, key=key)
-print(x2.shape)
-
-
 def _pad_kernel_to_target(kernel: Array, target_size: tuple[int, int]) -> Array:
     kh, kw = kernel.shape[-2], kernel.shape[-1]
     th, tw = target_size

{equimo-0.4.6a4 → equimo-0.4.7}/src/equimo/layers/sharing.py

@@ -1,5 +1,5 @@
 import math
-from typing import List, Optional
+from typing import List, Optional, Tuple
 
 import equinox as eqx
 import jax
@@ -72,8 +72,8 @@ class LayerSharing(eqx.Module):
 
     repeat: int = eqx.field(static=True)
 
-    loras: List[LoRA]
-    dropouts: List[eqx.nn.Dropout]
+    loras: Tuple[LoRA, ...]
+    dropouts: Tuple[eqx.nn.Dropout, ...]
     f: eqx.Module
 
     def __init__(
@@ -93,8 +93,8 @@ class LayerSharing(eqx.Module):
         keys = jr.split(key, repeat)
         self.repeat = repeat
 
-        self.dropouts = [eqx.nn.Dropout(drop_rate) for i in range(self.repeat)]
-        self.loras = [
+        self.dropouts = tuple(eqx.nn.Dropout(drop_rate) for i in range(self.repeat))
+        self.loras = tuple(
             LoRA(
                 in_features=dim,
                 out_features=dim,
@@ -103,7 +103,7 @@ class LayerSharing(eqx.Module):
                 key=keys[i],
             )
             for i in range(self.repeat)
-        ]
+        )
 
         self.f = f
 

{equimo-0.4.6a4 → equimo-0.4.7}/src/equimo/models/fastervit.py

@@ -1,4 +1,4 @@
-from typing import Callable, List, Literal, Optional
+from typing import Callable, List, Literal, Optional, Tuple
 
 import equinox as eqx
 import jax
@@ -157,7 +157,7 @@ class BlockChunk(eqx.Module):
     window_size: bool = eqx.field(static=True)
     do_gt: bool = eqx.field(static=True)
 
-    blocks: List[eqx.Module]
+    blocks: Tuple[eqx.Module, ...]
     downsample: eqx.Module
     global_tokenizer: Optional[TokenInitializer]
 
@@ -195,7 +195,7 @@ class BlockChunk(eqx.Module):
             k for k, v in kwargs.items() if isinstance(v, list) and len(v) == depth
         ]
 
-        self.blocks = []
+        blocks = []
         for i in range(depth):
             config = kwargs | {k: kwargs[k][i] for k in keys_to_spread}
 
@@ -210,7 +210,7 @@ class BlockChunk(eqx.Module):
                 }
 
             wrapper = LayerSharingWithCT if self.is_hat else LayerSharing
-            self.blocks.append(
+            blocks.append(
                 wrapper(
                     dim=kwargs.get("dim"),
                     f=block(**config, key=block_subkeys[i]),
@@ -218,6 +218,7 @@ class BlockChunk(eqx.Module):
                     key=block_subkeys[i],
                 ),
             )
+        self.blocks = tuple(blocks)
 
         self.downsample = downsampler(dim=kwargs.get("dim"), key=key_ds)
 
@@ -331,7 +332,7 @@ class FasterViT(eqx.Module):
     """
 
     patch_embed: ConvPatchEmbed
-    blocks: List[eqx.Module]
+    blocks: Tuple[eqx.Module, ...]
     norm: eqx.Module
     head: eqx.Module
 
@@ -398,7 +399,7 @@ class FasterViT(eqx.Module):
         hat = to_list(hat, n_chunks)
         attn_layer = to_list(attn_layer, n_chunks)
         window_size = to_list(window_size, n_chunks)
-        self.blocks = [
+        self.blocks = tuple(
             BlockChunk(
                 block=ConvBlock if i < 2 else HATBlock,
                 repeat=repeat,
@@ -427,7 +428,7 @@ class FasterViT(eqx.Module):
                 key=block_subkeys[i],
             )
             for i, depth in enumerate(depths)
-        ]
+        )
 
         num_features = int(dim * 2 ** (len(depths) - 1))
         self.norm = norm_layer(num_features)

{equimo-0.4.6a4 → equimo-0.4.7}/src/equimo/models/lowformer.py

@@ -15,7 +15,7 @@ from equimo.layers.norm import get_norm
 
 class BlockChunk(eqx.Module):
     residuals: list[bool] = eqx.field(static=True)
-    blocks: list[DSConv | MBConv | LowFormerBlock]
+    blocks: Tuple[DSConv | MBConv | LowFormerBlock, ...]
 
     def __init__(
         self,
@@ -123,7 +123,7 @@ class BlockChunk(eqx.Module):
                     )
                 residuals.append(False)
 
-        self.blocks = blocks
+        self.blocks = tuple(blocks)
         self.residuals = residuals
 
     def __call__(
@@ -146,7 +146,7 @@ class BlockChunk(eqx.Module):
 
 class LowFormer(eqx.Module):
     input_stem: eqx.nn.Sequential
-    blocks: list[BlockChunk]
+    blocks: Tuple[BlockChunk, ...]
     head: eqx.nn.Linear | eqx.nn.Identity
 
     def __init__(
@@ -220,7 +220,7 @@ class LowFormer(eqx.Module):
             ]
         )
 
-        self.blocks = [
+        self.blocks = tuple(
             BlockChunk(
                 in_channels=widths[i - 1] if i > 0 else width_stem,
                 out_channels=widths[i],
@@ -241,7 +241,7 @@ class LowFormer(eqx.Module):
             for i, (depth, att_stride, block_type, key_block) in enumerate(
                 zip(depths, att_strides, block_types, key_blocks)
             )
-        ]
+        )
 
         self.head = (
             eqx.nn.Linear(

{equimo-0.4.6a4 → equimo-0.4.7}/src/equimo/models/mlla.py

@@ -1,4 +1,4 @@
-from typing import List, Optional
+from typing import List, Optional, Tuple
 
 import equinox as eqx
 import jax
@@ -38,7 +38,7 @@ class Mlla(eqx.Module):
 
     patch_embed: eqx.Module
     pos_drop: eqx.Module
-    blocks: List[eqx.Module]
+    blocks: Tuple[eqx.Module, ...]
     head: eqx.Module
 
     def __init__(
@@ -52,7 +52,7 @@ class Mlla(eqx.Module):
         patch_size: int = 4,
         depths: List[int] = [2, 2, 6, 2],
         num_heads: List[int] = [3, 6, 12, 24],
-        attentions_layers: List[eqx.Module] | eqx.Module = LinearAttention,
+        attentions_layers: Tuple[eqx.Module, ...] | eqx.Module = LinearAttention,
         drop_rate: float = 0.0,
         drop_path_rate: float = 0.0,
         drop_path_uniform: bool = False,
@@ -101,8 +101,8 @@ class Mlla(eqx.Module):
         )
 
         num_heads = to_list(num_heads, n_chunks)
-        attentions_layers = to_list(attentions_layers, n_chunks)
-        self.blocks = [
+        attentions_layers = tuple(to_list(attentions_layers, n_chunks))
+        self.blocks = tuple(
             BlockChunk(
                 block=MllaBlock,
                 repeat=repeat,
@@ -124,7 +124,7 @@ class Mlla(eqx.Module):
                 key=block_subkeys[i],
             )
             for i, depth in enumerate(depths)
-        ]
+        )
 
         self.norm = eqx.nn.LayerNorm(self.num_features)
         self.head = (

{equimo-0.4.6a4 → equimo-0.4.7}/src/equimo/models/partialformer.py

@@ -98,7 +98,7 @@ class BlockChunk(eqx.Module):
     downsampler_contains_dropout: bool = eqx.field(static=True)
 
     posemb: eqx.Module
-    blocks: List[eqx.Module]
+    blocks: Tuple[eqx.Module, ...]
     downsample: eqx.Module
     qa_proj: eqx.Module
     qa_drop: eqx.Module
@@ -247,7 +247,7 @@ class PartialFormer(eqx.Module):
     qa_token: jnp.ndarray
     patch_embed: Stem
     pos_drop: eqx.nn.Dropout
-    blocks: List[eqx.Module]
+    blocks: Tuple[eqx.Module, ...]
     norm: eqx.Module
     head: eqx.Module
 

{equimo-0.4.6a4 → equimo-0.4.7}/src/equimo/models/reduceformer.py

@@ -14,7 +14,7 @@ from equimo.layers.norm import get_norm
 
 
 class BlockChunk(eqx.Module):
-    blocks: list[DSConv | MBConv | RFAttentionBlock]
+    blocks: Tuple[DSConv | MBConv | RFAttentionBlock, ...]
 
     def __init__(
         self,
@@ -126,7 +126,7 @@ class BlockChunk(eqx.Module):
                         )
                     )
 
-        self.blocks = blocks
+        self.blocks = tuple(blocks)
 
     def __call__(
         self,
@@ -148,7 +148,7 @@ class BlockChunk(eqx.Module):
 class ReduceFormer(eqx.Module):
     conv_stem: SingleConvBlock
     block_stem: BlockChunk
-    blocks: list[BlockChunk]
+    blocks: Tuple[BlockChunk, ...]
     head: eqx.nn.Linear | eqx.nn.Identity
 
     def __init__(
@@ -220,7 +220,7 @@ class ReduceFormer(eqx.Module):
             key=key_block_stem,
         )
 
-        self.blocks = [
+        self.blocks = tuple(
             BlockChunk(
                 in_channels=widths[i - 1] if i > 0 else width_stem,
                 out_channels=widths[i],
@@ -240,7 +240,7 @@ class ReduceFormer(eqx.Module):
             for i, (depth, block_type, key_block) in enumerate(
                 zip(depths, block_types, key_blocks)
             )
-        ]
+        )
 
         self.head = (
             eqx.nn.Linear(

{equimo-0.4.6a4 → equimo-0.4.7}/src/equimo/models/shvit.py

@@ -1,4 +1,4 @@
-from typing import List, Optional
+from typing import List, Optional, Tuple
 
 import equinox as eqx
 import jax
@@ -240,7 +240,7 @@ class SHViT(eqx.Module):
     """
 
     patch_embed: eqx.nn.Sequential
-    blocks: List[eqx.Module]
+    blocks: Tuple[eqx.Module, ...]
     head: eqx.Module
 
     def __init__(
@@ -316,7 +316,7 @@ class SHViT(eqx.Module):
         qk_dims = to_list(qk_dim, n_chunks)
         pdims = to_list(pdim, n_chunks)
         block_types = to_list(block_type, n_chunks)
-        self.blocks = [
+        self.blocks = tuple(
             BlockChunk(
                 block=BasicBlock,
                 repeat=repeat,
@@ -334,7 +334,7 @@ class SHViT(eqx.Module):
                 key=block_subkeys[i],
             )
             for i, depth in enumerate(depths)
-        ]
+        )
 
         self.norm = eqx.nn.LayerNorm(dims[-1])
         self.head = (

{equimo-0.4.6a4 → equimo-0.4.7}/src/equimo/models/vit.py

@@ -1,4 +1,4 @@
-from typing import Callable, List, Literal, Optional
+from typing import Callable, Literal, Optional, Tuple
 
 import equinox as eqx
 import jax
@@ -39,7 +39,7 @@ class BlockChunk(eqx.Module):
     downsampler_contains_dropout: bool = eqx.field(static=True)
 
     posemb: eqx.Module
-    blocks: List[eqx.Module]
+    blocks: Tuple[eqx.Module]
     downsample: eqx.Module
 
     def __init__(
@@ -84,7 +84,7 @@ class BlockChunk(eqx.Module):
         for i in range(depth):
             config = kwargs | {k: kwargs[k][i] for k in keys_to_spread}
             blocks.append(block(**config, key=block_subkeys[i]))
-        self.blocks = blocks
+        self.blocks = tuple(blocks)
 
         self.downsample = downsampler(dim=dim, **downsampler_kwargs, key=key_ds)
 
@@ -145,7 +145,7 @@ class VisionTransformer(eqx.Module):
     cls_token: jax.Array | None
     reg_tokens: jax.Array | None
     mask_token: jax.Array | None
-    blocks: List[eqx.Module]
+    blocks: Tuple[eqx.Module, ...]
     pos_drop: eqx.nn.Dropout
     norm: eqx.Module
     local_cls_norm: eqx.Module | None
@@ -171,8 +171,8 @@ class VisionTransformer(eqx.Module):
         in_channels: int,
         dim: int,
         patch_size: int,
-        num_heads: int | List[int],
-        depths: List[int],
+        num_heads: int | list[int],
+        depths: list[int],
         *,
         key: PRNGKeyArray,
         use_mask_token: bool = False,
@@ -304,7 +304,7 @@ class VisionTransformer(eqx.Module):
         dims = to_list(dim, n_chunks)
         num_heads = to_list(num_heads, n_chunks)
         attn_layer = to_list(attn_layer, n_chunks)
-        self.blocks = [
+        self.blocks = tuple(
             BlockChunk(
                 block=block,
                 dim=dims[i],
@@ -328,7 +328,7 @@ class VisionTransformer(eqx.Module):
                 key=block_subkeys[i],
             )
             for i, depth in enumerate(depths)
-        ]
+        )
 
         self.norm = norm_layer(dim, eps=eps)
 

{equimo-0.4.6a4 → equimo-0.4.7}/src/equimo/models/vssd.py

@@ -1,4 +1,4 @@
-from typing import List, Optional
+from typing import List, Optional, Tuple
 
 import equinox as eqx
 import jax
@@ -43,7 +43,7 @@ class Vssd(eqx.Module):
 
     patch_embed: eqx.Module
     pos_drop: eqx.Module
-    blocks: List[eqx.Module]
+    blocks: Tuple[eqx.Module, ...]
     head: eqx.Module
 
     def __init__(
@@ -60,12 +60,12 @@ class Vssd(eqx.Module):
         patch_size: int = 4,
         depths: List[int] = [2, 4, 12, 4],
         num_heads: List[int] = [2, 4, 8, 16],
-        attentions_layers: List[eqx.Module] | eqx.Module = [
+        attentions_layers: Tuple[eqx.Module, ...] | eqx.Module = (
             Mamba2Mixer,
             Mamba2Mixer,
             Mamba2Mixer,
             Attention,
-        ],
+        ),
         drop_rate: float = 0.0,
         drop_path_rate: float = 0.0,
         drop_path_uniform: bool = False,
@@ -118,8 +118,8 @@ class Vssd(eqx.Module):
         )
 
         num_heads = to_list(num_heads, n_chunks)
-        attentions_layers = to_list(attentions_layers, n_chunks)
-        self.blocks = [
+        attentions_layers = tuple(to_list(attentions_layers, n_chunks))
+        self.blocks = tuple(
             BlockChunk(
                 block=MllaBlock,
                 repeat=repeat,
@@ -142,7 +142,7 @@ class Vssd(eqx.Module):
                 key=block_subkeys[i],
             )
             for i, depth in enumerate(depths)
-        ]
+        )
 
         self.norm = eqx.nn.LayerNorm(self.num_features)
         self.head = (

equimo-0.4.6a4/src/equimo/__init__.py

@@ -1 +0,0 @@
-__version__ = "0.4.6a4"