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

lalamo/safetensors.py

@@ -0,0 +1,97 @@
+import json
+import struct
+from collections.abc import Mapping
+from dataclasses import dataclass
+from io import BufferedReader, BufferedWriter
+from typing import Any, ClassVar, Self
+
+import cattrs
+import jax
+import jax.numpy as jnp
+import numpy as np
+from jaxtyping import Array
+
+from lalamo.utils import LazyDict
+
+SF2J = {
+    "BOOL": jnp.dtype(jnp.bool_),
+    "U8": jnp.dtype(jnp.uint8),
+    "I8": jnp.dtype(jnp.int8),
+    "I16": jnp.dtype(jnp.int16),
+    "U16": jnp.dtype(jnp.uint16),
+    "F16": jnp.dtype(jnp.float16),
+    "BF16": jnp.dtype(jnp.bfloat16),
+    "I32": jnp.dtype(jnp.int32),
+    "U32": jnp.dtype(jnp.uint32),
+    "F32": jnp.dtype(jnp.float32),
+    "C64": jnp.dtype(jnp.complex64),
+    "F64": jnp.dtype(jnp.float64),
+    "I64": jnp.dtype(jnp.int64),
+    "U64": jnp.dtype(jnp.uint64),
+}
+
+J2SF = {v: k for k, v in SF2J.items()}
+
+
+@dataclass(frozen=True)
+class SFTensorInfo:
+    _converter: ClassVar[cattrs.Converter] = cattrs.Converter()
+    _converter.register_structure_hook(jnp.dtype, lambda x, _: SF2J[x])
+    _converter.register_unstructure_hook(jnp.dtype, lambda x: J2SF[x])
+
+    dtype: jnp.dtype
+    shape: tuple[int, ...]
+    data_offsets: tuple[int, int]
+
+    @property
+    def start(self) -> int:
+        return self.data_offsets[0]
+
+    @property
+    def end(self) -> int:
+        return self.data_offsets[1]
+
+    @property
+    def size(self) -> int:
+        return self.end - self.start
+
+    @classmethod
+    def from_dict(cls, obj: dict) -> Self:
+        return cls._converter.structure(obj, cls)
+
+    def to_dict(self) -> dict:
+        return self._converter.unstructure(self)
+
+
+def safe_read(fd: BufferedReader) -> tuple[dict[str, str] | None, LazyDict[str, Array]]:
+    header_size = struct.unpack("<Q", fd.read(8))[0]
+    header: dict[str, dict[str, Any]] = json.loads(fd.read(header_size))
+    metadata: dict[str, str] | None = header.pop("__metadata__", None)
+    data_offset = fd.tell()
+
+    def _load_tensor(key: str) -> Array:
+        info = SFTensorInfo.from_dict(header[key])
+        fd.seek(data_offset + info.start)
+        return jnp.asarray(np.fromfile(fd, info.dtype, info.size // info.dtype.itemsize)).reshape(info.shape)
+
+    lazy_tensors = LazyDict(set(header.keys()), _load_tensor)
+    return (metadata, lazy_tensors)
+
+
+def safe_write(fd: BufferedWriter, tensors: Mapping[str, Array]) -> None:
+    sorted_tensors = dict(sorted(tensors.items(), key=lambda x: (-x[1].dtype.alignment, x[0])))
+
+    header_dict = {}
+    offset = 0
+    for key, tensor in sorted_tensors.items():
+        assert offset % tensor.dtype.alignment == 0
+        header_dict[key] = SFTensorInfo(tensor.dtype, tensor.shape, (offset, offset + tensor.nbytes)).to_dict()
+        offset += tensor.nbytes
+
+    data_alignment = max(8, next((t.dtype.alignment for t in sorted_tensors.values()), 1))
+    header = json.dumps(header_dict).encode()
+    header += b" " * (-len(header) % data_alignment)
+    fd.write(struct.pack("<Q", len(header)) + header)
+
+    for tensor in sorted_tensors.values():
+        jax.device_get(tensor).tofile(fd)

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

@@ -9,6 +9,13 @@ import jax.numpy as jnp
 from lalamo.models import LanguageModel
 
 
+def get_default_device_memory() -> int | None:
+    memory_stats = jax.local_devices()[0].memory_stats()
+    if memory_stats is None or "bytes_limit" not in memory_stats:
+        return None
+    return memory_stats["bytes_limit"]
+
+
 def estimate_memory_from_batchsize(
     model: LanguageModel,
     max_input_length: int,
@@ -23,7 +30,7 @@ def estimate_memory_from_batchsize(
                 max_output_length=max_output_length,
                 num_top_logits_to_return=num_logits_per_token,
             ),
-            backend="cpu", # cuda backend tries to allocate in .compile() and ooms
+            backend="cpu",  # cuda backend tries to allocate in .compile() and ooms
         )
         .lower(
             model,
@@ -39,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/utils.py

@@ -9,21 +9,17 @@ from collections.abc import (
     Sequence,
     ValuesView,
 )
-from contextlib import contextmanager
 from dataclasses import dataclass
-from pathlib import Path
 from typing import overload
 
 import einops
 import jax.numpy as jnp
 from jaxtyping import Array
-from safetensors import safe_open
 
 __all__ = [
     "MapDictValues",
     "MapSequence",
     "jax_uint4_to_packed_uint8",
-    "open_safetensors",
     "process_chat_template",
 ]
 
@@ -45,15 +41,6 @@ class LazyDict[K, V](Mapping[K, V]):
         return len(self.stored_keys)
 
 
-@contextmanager
-def open_safetensors(filename: Path | str) -> Iterator[tuple[Mapping[str, Array], Mapping[str, str]]]:
-    with safe_open(filename, framework="flax") as safetensors_nonsense:
-        yield (
-            LazyDict(set(safetensors_nonsense.keys()), safetensors_nonsense.get_tensor),
-            safetensors_nonsense.metadata(),
-        )
-
-
 @dataclass(frozen=True)
 class MapIterable[OldT, NewT](Iterable[NewT]):
     map_func: Callable[[OldT], NewT]

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lalamo
-Version: 0.5.16
+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
@@ -19,7 +19,6 @@ Requires-Dist: rich>=14.0.0
 Requires-Dist: thefuzz>=0.22.1
 Requires-Dist: tokenizers>=0.21.2
 Requires-Dist: typer>=0.15.1
-Requires-Dist: safetensors>=0.6.2
 Requires-Dist: polars>=1.33.1
 Requires-Dist: xxhash>=3.5.0
 Provides-Extra: cpu

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

@@ -1,18 +1,20 @@
-lalamo/__init__.py,sha256=FjfGsBVSl14mNsDoFJEwXMRUq1-Kg_lessRzlJNG3KM,815
-lalamo/common.py,sha256=5NUFD26yQgOnEEk3LaQnce8n-VwJxILkEpFesHZhtQU,3820
-lalamo/main.py,sha256=GgUT7lT48-XQuAEH7qzsDKG8Lx9iBf-sYBIRhZL9q7E,23978
-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/sampling.py,sha256=g_dNiJyZrRqoQIiLid4cr6nRT9N5tSz3GtHr8Bt4n-E,3404
-lalamo/utils.py,sha256=QwATVXAeHBsQEDyt_31SHgxFphFVZYHpv3ZaklXks9Y,4585
+lalamo/registry_abc.py,sha256=qTikqviqqeseNzkjqoyQvL4dEWJYWzN0rI05T-JNTmo,2187
+lalamo/safetensors.py,sha256=kUiTSgx2zhfD1hxV_AA1DOLaKAKzjRd_vOYZCFf0em0,3048
+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
@@ -22,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=RLySCIkmGiA1IVZgLeemssMBMo4hMYMpmBjV0cRwBb4,6586
+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
@@ -52,28 +54,28 @@ lalamo/model_import/model_specs/qwen.py,sha256=HvN080ILpOwkqJbRLMqCa8Z8ImlLfTwiE
 lalamo/model_import/model_specs/reka.py,sha256=dOUYbEMMvovQdzQuBO_DCsjGI39syhoKCvnxLkNEDCw,423
 lalamo/models/__init__.py,sha256=Vn5PcvSqKppIchkSZwQVTn_GpRvOOzZVxo5PUeDl6N8,283
 lalamo/models/classifier.py,sha256=LvL54crCVi4HVSIXuoaSLB_5jtcx74GL7kgdy2Y16Zc,2094
-lalamo/models/common.py,sha256=PDteofGxjSBWYw_mPxbN1DTUba70aOURrAIjl13SSHc,2954
-lalamo/models/language_model.py,sha256=QPeVEyhutSze7fSNhvOvwSoYt24QMk-dtTJkos38amY,13465
+lalamo/models/common.py,sha256=uU6eCHtIqMeC_aRGVo09NdpAtvQ6RKSbm6pumVvL8pc,2943
+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
@@ -81,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=4D8dPZCWsrpORb7y8pQ6VsiIg1Cblvvxe6gXCoYtcD4,2530
+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.16.dist-info/licenses/LICENSE,sha256=diHRfjSEJHD1nnEeMIfMRCjR3UERf8bT3eseD6b1ayA,1072
-lalamo-0.5.16.dist-info/METADATA,sha256=dcs0vT9RULTxt4cxJJmfjP-4UJi7ZkrifXAaSMAgKeU,3147
-lalamo-0.5.16.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-lalamo-0.5.16.dist-info/entry_points.txt,sha256=qli7qTfnBk5WP10rOGXXEckHMtt-atJMDWd8jN89Uks,43
-lalamo-0.5.16.dist-info/top_level.txt,sha256=VHvWL5JN5XRG36NsN_MieJ7EwRihEOrEjyDaTdFJ-aI,7
-lalamo-0.5.16.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,,