lalamo 0.5.17__py3-none-any.whl → 0.6.0__py3-none-any.whl

lalamo/modules/token_mixers/attention.py

@@ -10,7 +10,7 @@ from jax import vmap
 from jaxtyping import Array, Bool, DTypeLike, Float, Int, PRNGKeyArray
 
 from lalamo.common import dummy_array
-from lalamo.modules.common import ParameterTree, PositionalEmbeddingSelector
+from lalamo.modules.common import ParameterTree, PositionalEmbeddingSelector, require_array, require_tree
 from lalamo.modules.linear import LinearBase, LinearConfig
 from lalamo.modules.normalization import Normalization, NormalizationConfig
 from lalamo.modules.rope import PositionalEmbeddings
@@ -433,33 +433,15 @@ class Attention(TokenMixerBase[AttentionConfig, KVCacheLayer]):
             result["sinks"] = self.sinks
         return result
 
-    def import_weights(
-        self,
-        weights: ParameterTree[Array],
-    ) -> Self:
+    def import_weights(self, weights: ParameterTree[Array]) -> Self:
         assert isinstance(weights, Mapping)
-        assert isinstance(weights["qkv_projection"], Mapping)
-        assert isinstance(weights["out_projection"], Mapping)
-        if self.query_norm is not None:
-            assert isinstance(weights["query_norm"], Mapping)
-            query_norm = self.query_norm.import_weights(weights["query_norm"])
-        else:
-            query_norm = None
-        if self.key_norm is not None:
-            assert isinstance(weights["key_norm"], Mapping)
-            key_norm = self.key_norm.import_weights(weights["key_norm"])
-        else:
-            key_norm = None
-        if self.sinks is not None:
-            assert isinstance(weights["sinks"], Array)
-            sinks = weights["sinks"]
-        else:
-            sinks = None
         return replace(
             self,
-            qkv_projection=self.qkv_projection.import_weights(weights["qkv_projection"]),
-            out_projection=self.out_projection.import_weights(weights["out_projection"]),
-            query_norm=query_norm,
-            key_norm=key_norm,
-            sinks=sinks,
+            qkv_projection=self.qkv_projection.import_weights(require_tree(weights["qkv_projection"])),
+            out_projection=self.out_projection.import_weights(require_tree(weights["out_projection"])),
+            query_norm=self.query_norm.import_weights(require_tree(weights["query_norm"]))
+            if self.query_norm
+            else None,
+            key_norm=self.key_norm.import_weights(require_tree(weights["key_norm"])) if self.key_norm else None,
+            sinks=require_array(weights["sinks"]) if self.sinks is not None else None,
         )

lalamo/modules/token_mixers/mamba.py

@@ -10,7 +10,7 @@ from einops import einsum, rearrange
 from jax import vmap
 from jaxtyping import Array, DTypeLike, Float, Int, PRNGKeyArray
 
-from lalamo.common import ParameterTree, dummy_array
+from lalamo.common import ParameterTree, dummy_array, require_array, require_tree
 from lalamo.modules.activations import Activation
 from lalamo.modules.common import LalamoModule, PositionalEmbeddingSelector
 from lalamo.modules.linear import LinearBase, LinearConfig
@@ -149,16 +149,10 @@ class SeparableCausalConv(LalamoModule[SeparableCausalConvConfig]):
 
     def import_weights(self, weights: ParameterTree[Array]) -> "SeparableCausalConv":
         assert isinstance(weights, Mapping)
-        assert isinstance(weights["weights"], Array)
-        if self.biases is not None:
-            assert isinstance(weights["biases"], Array)
-            biases = weights["biases"]
-        else:
-            biases = None
         return replace(
             self,
-            weights=weights["weights"],
-            biases=biases,
+            weights=require_array(weights["weights"]),
+            biases=require_array(weights["biases"]) if self.biases is not None else None,
         )
 
 
@@ -532,22 +526,13 @@ class Mamba2(TokenMixerBase[Mamba2Config, Mamba2StateLayer]):
             "gate_bias": self.gate_bias,
         }
 
-    def import_weights(
-        self,
-        weights: ParameterTree[Array],
-    ) -> Self:
+    def import_weights(self, weights: ParameterTree[Array]) -> Self:
         assert isinstance(weights, Mapping)
-        assert isinstance(weights["in_projection"], Mapping)
-        assert isinstance(weights["out_projection"], Mapping)
-        assert isinstance(weights["conv"], Mapping)
-        assert isinstance(weights["skip_connection_weight"], Array)
-        assert isinstance(weights["gate_bias"], Array)
-
         return replace(
             self,
-            in_projection=self.in_projection.import_weights(weights["in_projection"]),
-            out_projection=self.out_projection.import_weights(weights["out_projection"]),
-            conv=self.conv.import_weights(weights["conv"]),
-            skip_connection_weight=weights["skip_connection_weight"],
-            gate_bias=weights["gate_bias"],
+            in_projection=self.in_projection.import_weights(require_tree(weights["in_projection"])),
+            out_projection=self.out_projection.import_weights(require_tree(weights["out_projection"])),
+            conv=self.conv.import_weights(require_tree(weights["conv"])),
+            skip_connection_weight=require_array(weights["skip_connection_weight"]),
+            gate_bias=require_array(weights["gate_bias"]),
         )

lalamo/modules/token_mixers/short_conv.py

@@ -6,7 +6,7 @@ import equinox as eqx
 from jax import vmap
 from jaxtyping import Array, DTypeLike, Float, Int, PRNGKeyArray
 
-from lalamo.common import ParameterTree
+from lalamo.common import ParameterTree, require_tree
 from lalamo.modules.common import PositionalEmbeddingSelector
 from lalamo.modules.linear import LinearBase, LinearConfig
 from lalamo.modules.rope import PositionalEmbeddings
@@ -151,18 +151,11 @@ class ShortConv(TokenMixerBase[ShortConvConfig, ShortConvStateLayer]):
             "out_projection": self.out_projection.export_weights(),
         }
 
-    def import_weights(
-        self,
-        weights: ParameterTree[Array],
-    ) -> Self:
+    def import_weights(self, weights: ParameterTree[Array]) -> Self:
         assert isinstance(weights, Mapping)
-        assert isinstance(weights["in_projection"], Mapping)
-        assert isinstance(weights["conv"], Mapping)
-        assert isinstance(weights["out_projection"], Mapping)
-
         return replace(
             self,
-            in_projection=self.in_projection.import_weights(weights["in_projection"]),
-            conv=self.conv.import_weights(weights["conv"]),
-            out_projection=self.out_projection.import_weights(weights["out_projection"]),
+            in_projection=self.in_projection.import_weights(require_tree(weights["in_projection"])),
+            conv=self.conv.import_weights(require_tree(weights["conv"])),
+            out_projection=self.out_projection.import_weights(require_tree(weights["out_projection"])),
         )

lalamo/modules/transformer.py

@@ -7,7 +7,7 @@ import jax
 from jax import vmap
 from jaxtyping import Array, DTypeLike, Float, Int, PRNGKeyArray
 
-from lalamo.common import ParameterTree
+from lalamo.common import ParameterTree, require_tree
 from lalamo.modules.token_mixers import AttentionConfig
 from lalamo.modules.utils import vmap_twice
 
@@ -182,7 +182,8 @@ class Transformer(LalamoModule[TransformerConfig]):
     ) -> TransformerResult:
         if inner_features.ndim != 3:
             raise ValueError(
-                f"inner_features must be a 3D array of size (batch_size, sequence_length, hidden_dim), got {inner_features.shape}",
+                "inner_features must be a 3D array of size (batch_size, sequence_length, hidden_dim),"
+                f" got {inner_features.shape}",
             )
         if token_positions.ndim != 2:
             raise ValueError(
@@ -251,35 +252,24 @@ class Transformer(LalamoModule[TransformerConfig]):
             result["local_rope"] = self.local_rope.export_weights()
         return result
 
-    def import_weights(
-        self,
-        weights: ParameterTree[Array],
-    ) -> Self:
+    def import_weights(self, weights: ParameterTree[Array]) -> Self:
         assert isinstance(weights, Mapping)
         assert isinstance(weights["layers"], Sequence)
-        assert isinstance(weights["output_norm"], Mapping)
-
         if self.global_rope:
-            assert isinstance(weights["global_rope"], Mapping)
-            global_rope = self.global_rope.import_weights(weights["global_rope"])
+            global_rope = self.global_rope.import_weights(require_tree(weights["global_rope"]))
         else:
             global_rope = None
-
         if self.local_rope:
-            assert isinstance(weights["local_rope"], Mapping)
-            local_rope = self.local_rope.import_weights(weights["local_rope"])
+            local_rope = self.local_rope.import_weights(require_tree(weights["local_rope"]))
         else:
             local_rope = None
-
-        layers = []
-        for layer, layer_weights in zip(self.layers, weights["layers"], strict=True):
-            assert isinstance(layer_weights, Mapping)
-            layers.append(layer.import_weights(layer_weights))
-
+        layers = [
+            layer.import_weights(require_tree(lw)) for layer, lw in zip(self.layers, weights["layers"], strict=True)
+        ]
         return replace(
             self,
             global_rope=global_rope,
             layers=tuple(layers),
-            output_norm=self.output_norm.import_weights(weights["output_norm"]),
+            output_norm=self.output_norm.import_weights(require_tree(weights["output_norm"])),
             local_rope=local_rope,
         )

lalamo/modules/transformer_layer.py

@@ -9,7 +9,7 @@ import jax.numpy as jnp
 from jax import vmap
 from jaxtyping import Array, DTypeLike, Float, Int, PRNGKeyArray
 
-from lalamo.common import ParameterTree
+from lalamo.common import ParameterTree, require_tree
 
 from .common import ForwardPassMode, LalamoModule, PositionalEmbeddingSelector
 from .mlp import MLPBase, MLPConfig, MLPForwardPassConfig
@@ -293,38 +293,26 @@ class TransformerLayer(LalamoModule[TransformerLayerConfig]):
             result["post_mlp_norm"] = self.post_mlp_norm.export_weights()
         return result
 
-    def import_weights(
-        self,
-        weights: ParameterTree[Array],
-    ) -> Self:
+    def import_weights(self, weights: ParameterTree[Array]) -> Self:
         assert isinstance(weights, Mapping)
-        assert isinstance(weights["mixer"], Mapping)
-        assert isinstance(weights["mlp"], Mapping)
-        assert isinstance(weights["pre_mlp_norm"], Mapping)
-
         if self.post_mixer_norm is not None:
-            assert isinstance(weights["post_mixer_norm"], Mapping)
-            post_mixer_norm = self.post_mixer_norm.import_weights(
-                weights["post_mixer_norm"],
-            )
+            post_mixer_norm = self.post_mixer_norm.import_weights(require_tree(weights["post_mixer_norm"]))
         else:
             post_mixer_norm = None
         if self.post_mlp_norm is not None:
-            assert isinstance(weights["post_mlp_norm"], Mapping)
-            post_mlp_norm = self.post_mlp_norm.import_weights(weights["post_mlp_norm"])
+            post_mlp_norm = self.post_mlp_norm.import_weights(require_tree(weights["post_mlp_norm"]))
         else:
             post_mlp_norm = None
         if self.pre_mixer_norm is not None:
-            assert isinstance(weights["pre_mixer_norm"], Mapping)
-            pre_mixer_norm = self.pre_mixer_norm.import_weights(weights["pre_mixer_norm"])
+            pre_mixer_norm = self.pre_mixer_norm.import_weights(require_tree(weights["pre_mixer_norm"]))
         else:
             pre_mixer_norm = None
         return replace(
             self,
             pre_mixer_norm=pre_mixer_norm,
-            mixer=self.mixer.import_weights(weights["mixer"]),
+            mixer=self.mixer.import_weights(require_tree(weights["mixer"])),
             post_mixer_norm=post_mixer_norm,
-            pre_mlp_norm=self.pre_mlp_norm.import_weights(weights["pre_mlp_norm"]),
-            mlp=self.mlp.import_weights(weights["mlp"]),
+            pre_mlp_norm=self.pre_mlp_norm.import_weights(require_tree(weights["pre_mlp_norm"])),
+            mlp=self.mlp.import_weights(require_tree(weights["mlp"])),
             post_mlp_norm=post_mlp_norm,
         )

lalamo/registry_abc.py

@@ -1,5 +1,5 @@
 from abc import ABC, ABCMeta
-from typing import Any
+from typing import Any, Self
 from weakref import WeakSet
 
 __all__ = ["RegistryABC", "RegistryMeta"]
@@ -29,7 +29,7 @@ class RegistryMeta(ABCMeta):
 
         # Detect and remember the root exactly once
         if RegistryMeta._ROOT is None and name == "RegistryABC":
-            RegistryMeta._ROOT = cls  # type: ignore[assignment]
+            RegistryMeta._ROOT = cls
             return
 
         root = RegistryMeta._ROOT
@@ -58,6 +58,6 @@ class RegistryABC(ABC, metaclass=RegistryMeta):
     """
 
     @classmethod
-    def __descendants__(cls) -> tuple[type, ...]:
-        reg: WeakSet[type] = getattr(cls, RegistryMeta._REG_ATTR)  # noqa: SLF001
+    def __descendants__(cls) -> tuple[type[Self], ...]:
+        reg: WeakSet[type[Self]] = getattr(cls, RegistryMeta._REG_ATTR)  # noqa: SLF001
         return tuple(reg)

lalamo/sampling.py

@@ -1,5 +1,6 @@
 from abc import abstractmethod
 from collections.abc import Iterable
+from math import log
 
 import equinox as eqx
 import jax
@@ -10,6 +11,7 @@ __all__ = [
     "BanTokensPolicy",
     "CompositePolicy",
     "GreedyPolicy",
+    "MinPPolicy",
     "SamplingPolicy",
     "TemperaturePolicy",
     "TopKPolicy",
@@ -64,6 +66,15 @@ class TopPPolicy(SamplingPolicy):
         return jnp.where(to_remove_unsorted, -jnp.inf, logits)
 
 
+class MinPPolicy(SamplingPolicy):
+    p: float = eqx.field(static=True)
+
+    def process_logits(self, logits: Float[Array, " vocabulary"]) -> Float[Array, " vocabulary"]:
+        max_logit = jnp.max(logits)
+        logit_cutoff = max_logit + log(self.p)
+        return jnp.where(logits >= logit_cutoff, logits, -jnp.inf)
+
+
 class BanTokensPolicy(SamplingPolicy):
     banned_tokens: tuple[int, ...] = eqx.field(static=True)
 
@@ -85,6 +96,7 @@ def make_policy(
     temperature: float | None = None,
     top_k: int | None = None,
     top_p: float | None = None,
+    min_p: float | None = None,
     banned_tokens: Iterable[int] | None = None,
 ) -> SamplingPolicy:
     policies = []
@@ -96,4 +108,6 @@ def make_policy(
         policies.append(TopKPolicy(top_k))
     if top_p is not None:
         policies.append(TopPPolicy(top_p))
+    if min_p is not None:
+        policies.append(MinPPolicy(min_p))
     return CompositePolicy(tuple(policies))

lalamo/speculator/estimator.py

@@ -46,9 +46,9 @@ def estimate_memory_from_batchsize(
     assert hasattr(memory_analysis, "temp_size_in_bytes")
 
     return (
-        memory_analysis.argument_size_in_bytes  # type: ignore (pyright bug)
-        + memory_analysis.output_size_in_bytes  # type: ignore (pyright bug)
-        + memory_analysis.temp_size_in_bytes  # type: ignore (pyright bug)
+        memory_analysis.argument_size_in_bytes
+        + memory_analysis.output_size_in_bytes
+        + memory_analysis.temp_size_in_bytes
     )
 
 

lalamo/speculator/ngram.py

@@ -129,7 +129,7 @@ class NGramSpeculator(Speculator):
 
         return (
             memoryview(self.ngram_keys)[idx_start:idx_end],
-            memoryview(self.ngram_values)[idx_start:idx_end],  # type: ignore (typechecker bug)
+            memoryview(self.ngram_values)[idx_start:idx_end].cast("f"), # noop cast to make typechecker happy
             memoryview(self.ngram_counts)[seq_hash : (seq_hash + 1)],
         )
 

{lalamo-0.5.17.dist-info → lalamo-0.6.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lalamo
-Version: 0.5.17
+Version: 0.6.0
 Summary: JAX library for optimization and export of models for use with the UZU inference engine.
 Requires-Python: <4,>=3.12
 Description-Content-Type: text/markdown

{lalamo-0.5.17.dist-info → lalamo-0.6.0.dist-info}/RECORD

@@ -1,20 +1,20 @@
-lalamo/__init__.py,sha256=asVMPmQ7BUt7bYlcuNZ7SnOSJDJUiN9QhlU5lRUehSo,1387
-lalamo/commands.py,sha256=rU9T8Mx6s7itpk-dj5ToQ4PUpGPfdmmKlrF02l2kIS0,9967
-lalamo/common.py,sha256=5NUFD26yQgOnEEk3LaQnce8n-VwJxILkEpFesHZhtQU,3820
-lalamo/main.py,sha256=dE7Us9L6sfz9bp5rUSzGHUkG0Uon4xdju9dGGtXidZI,23888
-lalamo/message_processor.py,sha256=bSUAQg7CemLTnBV4LtPxJBicAalruDCA-JXjkTYPZ8U,5797
+lalamo/__init__.py,sha256=vJX2fu1UJdIdtVlBjxWtPNTw8dEge9G2PI1eQOGF-qA,1386
+lalamo/commands.py,sha256=zXyyrLTHhP9wouwtpX4RUZeEF6No-_9ee-y_GWGhw7k,10972
+lalamo/common.py,sha256=WaNJx20eUX4CBF50aym9lniGAiX-SzBJzDzO5Jh6zXA,4312
+lalamo/main.py,sha256=Tez84CtMxUi1ySuRSqQElu4Zr1UWs_Gw6HX1xtCZknQ,27383
+lalamo/message_processor.py,sha256=PMKte9YijT3h9N7DjTNp8H4V45A_qlDqJaubqFevLX8,5924
 lalamo/quantization.py,sha256=8o6ryIZLzzDYQuvBTboPfaVVdfijAKGpTxOcg3GKVD8,2752
-lalamo/registry_abc.py,sha256=ENjXiD_wEH100fNjG-W5Em1L_EQ0Lf0pdRhRGvf3qZk,2197
+lalamo/registry_abc.py,sha256=qTikqviqqeseNzkjqoyQvL4dEWJYWzN0rI05T-JNTmo,2187
 lalamo/safetensors.py,sha256=kUiTSgx2zhfD1hxV_AA1DOLaKAKzjRd_vOYZCFf0em0,3048
-lalamo/sampling.py,sha256=g_dNiJyZrRqoQIiLid4cr6nRT9N5tSz3GtHr8Bt4n-E,3404
+lalamo/sampling.py,sha256=GE6Av7zS-pr5Bg7FtOivRce7I0JIYuNYqfqsRe-yjQk,3867
 lalamo/utils.py,sha256=c88IP110gHZJ6hYDq7p36A9u-vLRM_YdavFom56gsNQ,4111
 lalamo/data/__init__.py,sha256=exfhBLxHrg7BWutM0tAln5QuIWlNQmOhaG2noFYxfPI,189
 lalamo/data/huggingface_message.py,sha256=-7lN9eIcETQzt1Pnx3d4d8p3_I7WYMNf4mp1P91N7fI,1115
 lalamo/data/lalamo_completions.py,sha256=U_m3UNSJASUFz3rJq_taZOtL_U4B8Oj-ndkTF-JH-v4,1509
 lalamo/data/utils.py,sha256=B96gLaULyStKYuR8wjFdTpFc6YIDC8EEvGh1eiMe_Ec,338
 lalamo/model_import/__init__.py,sha256=Z8pS9rbKKx1QgUy7KZtHxiNWlZhII3mdovT9d37vAxg,168
-lalamo/model_import/common.py,sha256=wvyGD-iLut_Pm3HjDMI05upqdtCW3HWeoeB0YmiFeqk,12419
-lalamo/model_import/huggingface_generation_config.py,sha256=mot6VQ6ezCtEhN6VjhnvaU-nR5P5T2BuBUgpFNnWJxU,1495
+lalamo/model_import/common.py,sha256=MIbvK3mxgrDSXea6jujvCOu9Jjyip6MXeTsJjNTBJAU,12325
+lalamo/model_import/huggingface_generation_config.py,sha256=xicv_kJOfIGlz4gi5fRFIkiAZ9_QRDLRtW8nKMm5tVU,2022
 lalamo/model_import/huggingface_tokenizer_config.py,sha256=xvwdmio7b9nhn2H3uMBVligiYj58JaCFCvHY3-8dBvM,2502
 lalamo/model_import/decoder_configs/__init__.py,sha256=YvlSsJqNEQPCNKcUzCw0MLjt8H3vcfjc4sz1OK7qdIQ,679
 lalamo/model_import/decoder_configs/common.py,sha256=L8PCgF5fIt3RqPlmLiJpBzDguKk9iTjk4XSItxwVG4c,3260
@@ -24,26 +24,26 @@ lalamo/model_import/decoder_configs/huggingface/common.py,sha256=YYIDEQy8x7lqL2q
 lalamo/model_import/decoder_configs/huggingface/gemma2.py,sha256=g8LH_GlSNyL04WWi596zI0rWsD3ahnfNjDk-9zZNcDE,4759
 lalamo/model_import/decoder_configs/huggingface/gemma3.py,sha256=UXiEyNqlD0Czc5Gj3n4hNqNDp9Ml5YzH1XZ6BXj0mgU,10223
 lalamo/model_import/decoder_configs/huggingface/gpt_oss.py,sha256=MBCoPbuWyzbJiBRtHOtpaPHJjQ1UVCAYcVrfIejTnlQ,7446
-lalamo/model_import/decoder_configs/huggingface/lfm2.py,sha256=vrBMxtiKEg0eHNDL_bWM9odlrsab7jlMXEY8vjEB7-c,7595
+lalamo/model_import/decoder_configs/huggingface/lfm2.py,sha256=tOx4EsDGRd-87E1Q94DkbGlRBeIvBOvapfr9WeUxFYE,8027
 lalamo/model_import/decoder_configs/huggingface/llama.py,sha256=pGuBQTY6qpx6CriWwdsLpuTSRS7ECoTP1kt5pSKRlNQ,8549
-lalamo/model_import/decoder_configs/huggingface/llamba.py,sha256=ANB-vQK8U-zVFubZSTDXXt2S70T5SVOGzf7eOVvPzIQ,5773
+lalamo/model_import/decoder_configs/huggingface/llamba.py,sha256=NVvr7_3bfcLHGRrHG3b0IylgTt-knH31oLz3yFqrkqQ,5775
 lalamo/model_import/decoder_configs/huggingface/mistral.py,sha256=MDGC0ivzJuUpOC11n8vFdcVzqccUyaRw_hkL74mVlAg,4599
-lalamo/model_import/decoder_configs/huggingface/modern_bert.py,sha256=A8nNIMhPVumvPWIFR3RexRc6XkFyUd_3mmNpmvyPEGE,8816
+lalamo/model_import/decoder_configs/huggingface/modern_bert.py,sha256=Crh20pjSa35fP22D3J-29mv4yzdrjzW6VhOjY4Tasmg,8801
 lalamo/model_import/decoder_configs/huggingface/qwen2.py,sha256=n3qIANMPbtQsTtk5QEWWFZ6R85eDxR_kaZd0NDlJ3T4,5786
 lalamo/model_import/decoder_configs/huggingface/qwen3.py,sha256=i99mfL2DbeJ0l5aFRV84MTT-PsWf6q-8B-SGPIVGe1w,7522
 lalamo/model_import/loaders/__init__.py,sha256=3THc1wQ4EPBzQkL_4EaKCa7Ev5Z7oczcvc4AHy9v5EI,228
 lalamo/model_import/loaders/common.py,sha256=kkugV-bMQlN1zvGHoj3uc7z0FbXKoMtXEBTvyu4KxK4,1844
-lalamo/model_import/loaders/executorch.py,sha256=t2Ey_mBMNC8bTSTdYWjuGXdPTRoohFlYrqtWyNkBU_8,9219
-lalamo/model_import/loaders/huggingface.py,sha256=qWdzoSvHvb_3prn2kwfxgnYPW2bVB0Q49m_wyRYha8Q,34677
+lalamo/model_import/loaders/executorch.py,sha256=JCeylxmkXT2iOfVmrvgAyP-9Th-96w3sRtssIW43Ag4,9187
+lalamo/model_import/loaders/huggingface.py,sha256=4zIKuYd5-BC1nkf6rtuKxnOmefEWafv6yXuKEdxg9p4,34629
 lalamo/model_import/loaders/utils.py,sha256=eiX3WKFRrAfBY-dugodscNInl5o5w3KmVcgma4atpGY,2456
 lalamo/model_import/model_specs/__init__.py,sha256=JISqwJkloQkGD2jvi1MakNEWapIwlNXXVi5giZyXB74,1275
-lalamo/model_import/model_specs/common.py,sha256=8ALKxHrt8uK4XiqjK25NwZj1CC7DM7jlYcFVZPGkFrw,6643
+lalamo/model_import/model_specs/common.py,sha256=OcE6wzDz4MsETxYdcOvRT6x6_NpsyBeIlK1Zl6qkMMo,6823
 lalamo/model_import/model_specs/deepseek.py,sha256=Umef93_ZBuq93yYsejIRNwj3udoln1gHfrv3SK5jyMo,417
 lalamo/model_import/model_specs/essential_ai.py,sha256=xbHcwRpAWhR9gOgypVzcgunFspoUEk3iNsw-46CVR4o,390
 lalamo/model_import/model_specs/gemma.py,sha256=dwKwOHU1sBJNLFAwtEyydsRUF9QENN3SHtjbfqtOSic,3876
 lalamo/model_import/model_specs/gpt_oss.py,sha256=PLo0QGrXKdX61ReTRdyOaP_EH3Dmj5lp3fpJjZRwRVA,542
 lalamo/model_import/model_specs/huggingface.py,sha256=TEkU8y95_hmUWyF-Q5hn0dE2SvXbApghAsQwhWRu4D0,431
-lalamo/model_import/model_specs/lfm2.py,sha256=uzuFbcj4Wj2OqL7XJE8Q431VYZelS_HkfPFpl7rJuJY,1038
+lalamo/model_import/model_specs/lfm2.py,sha256=wg4Ggt6BbMO4ScJ6h8tjvBc3IVSrMudESQxjleUF9Ds,2198
 lalamo/model_import/model_specs/llama.py,sha256=TxhKbIBFmGV2NopOg_k3ltsKlJccbxKyu-GQ7hYWCyw,3140
 lalamo/model_import/model_specs/llamba.py,sha256=Ic3sWTv34FLJ4fG6OR_Mc5goGJQR6fa5b2WbVXbn9FA,1471
 lalamo/model_import/model_specs/mirai.py,sha256=eifYVV5-fABiLH6rr82_DiVFtDyqpW0vbvXCYsQQzto,617
@@ -55,27 +55,27 @@ lalamo/model_import/model_specs/reka.py,sha256=dOUYbEMMvovQdzQuBO_DCsjGI39syhoKC
 lalamo/models/__init__.py,sha256=Vn5PcvSqKppIchkSZwQVTn_GpRvOOzZVxo5PUeDl6N8,283
 lalamo/models/classifier.py,sha256=LvL54crCVi4HVSIXuoaSLB_5jtcx74GL7kgdy2Y16Zc,2094
 lalamo/models/common.py,sha256=uU6eCHtIqMeC_aRGVo09NdpAtvQ6RKSbm6pumVvL8pc,2943
-lalamo/models/language_model.py,sha256=QPeVEyhutSze7fSNhvOvwSoYt24QMk-dtTJkos38amY,13465
+lalamo/models/language_model.py,sha256=HtFS-R4Uqr7SohFstoAZFVrJI293N9cG_LVkXhZxgFI,13546
 lalamo/modules/__init__.py,sha256=OHIQn08jx2c3L2KIQA-7SJ4yVb2E5m6T6FqTHFJTDdM,4006
-lalamo/modules/activations.py,sha256=U3qTQtZawPAUcoqbkIJnmTYcaNiQuSPMLcBeJ398GhI,1022
-lalamo/modules/classifier.py,sha256=_jtJ3INEq1dJP5HpUmcDk9YYzpRYlQ04zvFGaWBV6Lg,12101
-lalamo/modules/common.py,sha256=dqDEOi-C3H4U9iWUisU32RA-wRDCGuaUNGbObRBhyQM,3315
-lalamo/modules/decoder.py,sha256=Opd3QIq1mpGr9P7sLH-Fryitlfp6ESTpcX71vgm89t0,7129
-lalamo/modules/embedding.py,sha256=LLiH8mTu81JSpUTj-XhsrVIUfl_GhapnXxw1yGSUBgM,28428
-lalamo/modules/linear.py,sha256=XfIYhmpk-bwNHIzIgsL48ZUTclHD2KB4uXHMw9NTE-8,42991
-lalamo/modules/mlp.py,sha256=bL3sQ46vCNt1MBRwlzmXZx9nQfRe4axpGe5UOFVanBI,17959
-lalamo/modules/mlx_interop.py,sha256=FdfU_1iES-HQ9r4K0SkYwJTyvE0f-_T5ursNCjPLZKY,467
+lalamo/modules/activations.py,sha256=25F4XytJMIwPPmUbxiDUrcrdUi4c-O9SUbwv9lnZbuU,992
+lalamo/modules/classifier.py,sha256=Q5eNzJ68to6JGk8IDZiKv6Rmwh15UyT2xC52tP5njoQ,11767
+lalamo/modules/common.py,sha256=Rc9zenrUMntDKZydI1tzt1ZIY8ggfyk3ZDB-xi81ibw,3406
+lalamo/modules/decoder.py,sha256=I30fptNifcdw9OOCU50aZnEqsJ2X4VM9YXdtRkxbqGc,7014
+lalamo/modules/embedding.py,sha256=PdNy4tGt9F1zve4X73WKNS0DXL-nHUFOlZmGFUAarkQ,27727
+lalamo/modules/linear.py,sha256=4xIhmeouD7R10lt8KJBLxgypVXYhpGmXdHUc-96Upfk,42871
+lalamo/modules/mlp.py,sha256=ogxi9q8J38FnuBkAtC7_KTMc7JZG4BRdsAHYprHZNvM,17690
+lalamo/modules/mlx_interop.py,sha256=kgCm6cPvY62ZNY3icuyKY0bow50j73UdyfVym2WqEUk,483
 lalamo/modules/normalization.py,sha256=cBdOq6OmJssunVeEwFRJD0BDhgFAN7J8gOKwzIUAY8I,3005
-lalamo/modules/rope.py,sha256=rCik7vBNqRXYg3LGbmc1mezPRNbIYMg5cydTFpQy-eU,10157
+lalamo/modules/rope.py,sha256=HbIv5ESLGNAK47HAtqu1whLLUa20Sb28U8kEs6KclZM,10127
 lalamo/modules/torch_interop.py,sha256=-mujd1zI4ec2w92Hd50RtDa0K3jl6ZSnPxc5r3Fp9nU,916
-lalamo/modules/transformer.py,sha256=4olEO8Eh7U6RwSnaECn39ooPuTKUZp_6QmvO6vdirrQ,10532
-lalamo/modules/transformer_layer.py,sha256=ZYmGR2Ej328l7K-YpV4eEiBk8SzLsw1RiuSiUP94UpY,12731
+lalamo/modules/transformer.py,sha256=9FD2k_5qwDHYUG5_6M0wVI9-YxfMv0mXlHS-QKiKcP4,10319
+lalamo/modules/transformer_layer.py,sha256=mOqmfVpT7yfHpU87Koso3lvjH5zc-hgPvgVgk04r6ck,12412
 lalamo/modules/utils.py,sha256=t_TayWT6g5LtYKhJaod-u_COWaI_VbNd3eYek9Nj0lc,441
-lalamo/modules/token_mixers/__init__.py,sha256=z6x8cNjis6xIi_2llIoByKqMF2W4xJ05rDnxitHQ3jU,1139
-lalamo/modules/token_mixers/attention.py,sha256=gkGMFah2OHB_tyJpkshM1KhMnzG6U7Xt273MkBvDk58,16584
+lalamo/modules/token_mixers/__init__.py,sha256=lwxUl0eG5IvuVc_HOsINP2vtbv9F0cUmSNHFHaEmPGk,1109
+lalamo/modules/token_mixers/attention.py,sha256=ielw1-KWBfCPCPmzSHgM0TaSUcmSkWKTxrN3N_FsGm4,16144
 lalamo/modules/token_mixers/common.py,sha256=CcrbXXvGU27uxGLh5L-G8VDtcOiW5Wpm13uBEOd6lVg,1986
-lalamo/modules/token_mixers/mamba.py,sha256=fo8xvvmIQss2lKLhav19Jzk1-hTykNp2sjcN6ntcWj4,18789
-lalamo/modules/token_mixers/short_conv.py,sha256=93SmoVsuAtdX4ckAkvhHXHiO67pU6soYFpBZxdPFEwc,5219
+lalamo/modules/token_mixers/mamba.py,sha256=d5hq1M_fk467EKm9Nyu3xu5OyrUmp04eeTyPSKh4HVI,18422
+lalamo/modules/token_mixers/short_conv.py,sha256=xLoTiN44cFuPva9gdjjC8EQqI5RYSWRd_YFxTpyOsmE,5076
 lalamo/modules/token_mixers/state/__init__.py,sha256=OKWPmiwszMWgwamewoVHd28owanHAO2j2e30Iivtv-4,384
 lalamo/modules/token_mixers/state/common.py,sha256=dcwBevAdeJpBjf7_YRk7TKrJHsCnpljhfzZy-3h9898,661
 lalamo/modules/token_mixers/state/kv_cache.py,sha256=QfnS3XgSmyDI9MBUbeLI4ABHLxiMcXDbZsqe0fd3KQo,8788
@@ -83,13 +83,13 @@ lalamo/modules/token_mixers/state/mamba_state.py,sha256=LHzJvNE6MkB7nrsZSNto6pxb
 lalamo/modules/token_mixers/state/short_conv_state.py,sha256=osjcDHoeFWQaUoOROzeJe8F1qC8rvqunimGD4CuIDHo,895
 lalamo/speculator/__init__.py,sha256=9-tmZcbCom_lIGpJYn6xLlnEahFLFidpqmgkafmu--k,456
 lalamo/speculator/common.py,sha256=PudF_gkpe5_nQ-57sAC-foE1xCy_H2Axh5KwRoA86lo,587
-lalamo/speculator/estimator.py,sha256=j-zmhy3RxYDmQ7W0FMTmDk3i275r_Vg1s4NCaS4c_SQ,2760
+lalamo/speculator/estimator.py,sha256=S_TRwMnjWg5qt9le2AYua_Vmo6QkIT-0Si7TjCfC7xc,2670
 lalamo/speculator/inference.py,sha256=5GntUgj0HQLeLn3HIHnVX8EEO0EBzmKeP5-_U7kdFAM,3670
-lalamo/speculator/ngram.py,sha256=95mdfAWhx4d5XOnOwhyhElnvcy6nlUjYhcbJzqDs414,5875
+lalamo/speculator/ngram.py,sha256=Fy3A-oVxZql3gE5M5ot0hKPu0772-kcEPDvD9MkldpA,5889
 lalamo/speculator/utils.py,sha256=0wZoMMIzzk0Q-3zq5H5f-JBplePNHxywndkrNtOJOyo,1697
-lalamo-0.5.17.dist-info/licenses/LICENSE,sha256=diHRfjSEJHD1nnEeMIfMRCjR3UERf8bT3eseD6b1ayA,1072
-lalamo-0.5.17.dist-info/METADATA,sha256=16-W1J0wiwrmgMTgqiE9r3vxKRmZbGgZ-zS7bNACwTA,3113
-lalamo-0.5.17.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-lalamo-0.5.17.dist-info/entry_points.txt,sha256=qli7qTfnBk5WP10rOGXXEckHMtt-atJMDWd8jN89Uks,43
-lalamo-0.5.17.dist-info/top_level.txt,sha256=VHvWL5JN5XRG36NsN_MieJ7EwRihEOrEjyDaTdFJ-aI,7
-lalamo-0.5.17.dist-info/RECORD,,
+lalamo-0.6.0.dist-info/licenses/LICENSE,sha256=diHRfjSEJHD1nnEeMIfMRCjR3UERf8bT3eseD6b1ayA,1072
+lalamo-0.6.0.dist-info/METADATA,sha256=6cDrL3A2Hksel4WIH3kXyNx0YquoJy96siJDHX8MPa8,3112
+lalamo-0.6.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+lalamo-0.6.0.dist-info/entry_points.txt,sha256=qli7qTfnBk5WP10rOGXXEckHMtt-atJMDWd8jN89Uks,43
+lalamo-0.6.0.dist-info/top_level.txt,sha256=VHvWL5JN5XRG36NsN_MieJ7EwRihEOrEjyDaTdFJ-aI,7
+lalamo-0.6.0.dist-info/RECORD,,