lalamo 0.5.5__py3-none-any.whl → 0.5.7__py3-none-any.whl

lalamo/__init__.py

@@ -15,7 +15,7 @@ from lalamo.speculator import (
     SpeculatorTrainingEvent,
 )
 
-__version__ = "0.5.5"
+__version__ = "0.5.7"
 
 __all__ = [
     "AssistantMessage",

lalamo/models/common.py

@@ -75,7 +75,7 @@ class TextModel[ConfigT, ModelT: Decoder | Classifier](LalamoModule[ConfigT]):
         if messages is None:
             messages = [UserMessage("Tell me about London")]
 
-        token_ids = jnp.array(self.message_processor.tokenize_request(messages))[None:]
+        token_ids = jnp.array(self.message_processor.tokenize_request(messages))[None, :]
         _, num_tokens = token_ids.shape
         token_positions = jnp.arange(num_tokens)[None, :]
         return self.model(token_ids=token_ids, token_positions=token_positions, return_activation_trace=True)

lalamo/speculator/estimator.py

@@ -19,13 +19,14 @@ def estimate_memory_from_batchsize(
     memory_analysis = (
         jax.jit(
             functools.partial(
-                model.generate_tokens,
+                LanguageModel.generate_tokens,
                 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
         )
         .lower(
+            model,
             prompt_token_ids=jax.ShapeDtypeStruct((batch_size, max_input_length), jnp.int32),
             prompt_lengths_without_padding=jax.ShapeDtypeStruct((batch_size,), jnp.int32),
         )

lalamo/speculator/inference.py

@@ -30,12 +30,13 @@ def inference_collect_traces(
     generate_tokens_compiled = (
         jax.jit(
             functools.partial(
-                model.generate_tokens,
+                LanguageModel.generate_tokens,
                 max_output_length=max_output_length,
                 num_top_logits_to_return=num_top_logits_to_collect,
             ),
         )
         .lower(
+            model,
             prompt_token_ids=jax.ShapeDtypeStruct((batch_size, max_input_length), jnp.int32),
             prompt_lengths_without_padding=jax.ShapeDtypeStruct((batch_size,), jnp.int32),
         )
@@ -60,6 +61,7 @@ def inference_collect_traces(
         )
 
         generated = generate_tokens_compiled(
+            model,
             prompt_token_ids=padded,
             prompt_lengths_without_padding=length_without_padding,
         )

{lalamo-0.5.5.dist-info → lalamo-0.5.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lalamo
-Version: 0.5.5
+Version: 0.5.7
 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.5.dist-info → lalamo-0.5.7.dist-info}/RECORD

@@ -1,4 +1,4 @@
-lalamo/__init__.py,sha256=tkT6ErAGkHuPqJ6O32MPLbltRBh__kfmPM-SpU7iVko,762
+lalamo/__init__.py,sha256=ynqoSW4I6eQ92AR9jzr_8XMEiGpDGjMHHGUrjlTkmb0,762
 lalamo/common.py,sha256=5NUFD26yQgOnEEk3LaQnce8n-VwJxILkEpFesHZhtQU,3820
 lalamo/main.py,sha256=Lqv-lU6hRSqbJeVOheZoKllK1LaPHTuR_8jNTPC7UZg,23956
 lalamo/message_processor.py,sha256=bSUAQg7CemLTnBV4LtPxJBicAalruDCA-JXjkTYPZ8U,5797
@@ -48,7 +48,7 @@ lalamo/model_import/model_specs/polaris.py,sha256=Mw1-6bByjDmPIKlIUIV46CsmV5xUp_
 lalamo/model_import/model_specs/qwen.py,sha256=qzLmTveATmnwNFQSFJlffcXw7syFnrCmKf9ggkkkw1Y,7050
 lalamo/model_import/model_specs/reka.py,sha256=dOUYbEMMvovQdzQuBO_DCsjGI39syhoKCvnxLkNEDCw,423
 lalamo/models/__init__.py,sha256=HfAYc4mteKu7BMyIP8aTFliLKGWJzLk1kPcdKurb8eo,243
-lalamo/models/common.py,sha256=b2woUDgHhCUEowo6tEjQb4J7iRVQvrptyRQTYMKrLis,2952
+lalamo/models/common.py,sha256=PDteofGxjSBWYw_mPxbN1DTUba70aOURrAIjl13SSHc,2954
 lalamo/models/language_model.py,sha256=QPeVEyhutSze7fSNhvOvwSoYt24QMk-dtTJkos38amY,13465
 lalamo/models/router.py,sha256=7KZqHVhr2TA7Qh76KfwrvyfztfZnV-P-Ss11O8dzbRg,2013
 lalamo/modules/__init__.py,sha256=xWJ4OPAF4gKd0evYwXIK5kTnbH6nI55oLAePcoDDHQ0,3730
@@ -76,13 +76,13 @@ lalamo/modules/token_mixers/state/kv_cache.py,sha256=QfnS3XgSmyDI9MBUbeLI4ABHLxi
 lalamo/modules/token_mixers/state/mamba_state.py,sha256=LHzJvNE6MkB7nrsZSNto6pxbnMJCl--JOoe9Fkcc9Mg,1642
 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=SFF24TtxGVbb3hEziZ40A7aNdBd6_ZZJSX-CqmPhvA8,2503
-lalamo/speculator/inference.py,sha256=EzNLmdP-AOtIN8C39-fq-3ZnKX_yxOOwMufDEmi54zM,3624
+lalamo/speculator/estimator.py,sha256=4D8dPZCWsrpORb7y8pQ6VsiIg1Cblvvxe6gXCoYtcD4,2530
+lalamo/speculator/inference.py,sha256=5GntUgj0HQLeLn3HIHnVX8EEO0EBzmKeP5-_U7kdFAM,3670
 lalamo/speculator/ngram.py,sha256=95mdfAWhx4d5XOnOwhyhElnvcy6nlUjYhcbJzqDs414,5875
 lalamo/speculator/utils.py,sha256=0wZoMMIzzk0Q-3zq5H5f-JBplePNHxywndkrNtOJOyo,1697
-lalamo-0.5.5.dist-info/licenses/LICENSE,sha256=diHRfjSEJHD1nnEeMIfMRCjR3UERf8bT3eseD6b1ayA,1072
-lalamo-0.5.5.dist-info/METADATA,sha256=rv7bkPZ8-ZluxxZJHkGzT724Luws5wiMnGBAXNJbM58,3146
-lalamo-0.5.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-lalamo-0.5.5.dist-info/entry_points.txt,sha256=qli7qTfnBk5WP10rOGXXEckHMtt-atJMDWd8jN89Uks,43
-lalamo-0.5.5.dist-info/top_level.txt,sha256=VHvWL5JN5XRG36NsN_MieJ7EwRihEOrEjyDaTdFJ-aI,7
-lalamo-0.5.5.dist-info/RECORD,,
+lalamo-0.5.7.dist-info/licenses/LICENSE,sha256=diHRfjSEJHD1nnEeMIfMRCjR3UERf8bT3eseD6b1ayA,1072
+lalamo-0.5.7.dist-info/METADATA,sha256=ZzSQiR7KYoAXkBKOToa2bi9tCPRvVXvhnN5y6AO7wyQ,3146
+lalamo-0.5.7.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+lalamo-0.5.7.dist-info/entry_points.txt,sha256=qli7qTfnBk5WP10rOGXXEckHMtt-atJMDWd8jN89Uks,43
+lalamo-0.5.7.dist-info/top_level.txt,sha256=VHvWL5JN5XRG36NsN_MieJ7EwRihEOrEjyDaTdFJ-aI,7
+lalamo-0.5.7.dist-info/RECORD,,