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

lalamo/__init__.py

@@ -27,7 +27,7 @@ from lalamo.speculator import (
     SpeculatorTrainingEvent,
 )
 
-__version__ = "0.6.0"
+__version__ = "0.6.1"
 
 __all__ = [
     "AssistantMessage",

lalamo/modules/token_mixers/mamba.py

@@ -112,6 +112,7 @@ class SeparableCausalConv(LalamoModule[SeparableCausalConvConfig]):
     def __call__(
         self,
         inputs: Float[Array, "suffix_tokens channels"],
+        length_without_padding: Int[Array, ""] | int | None = None,
         state: Float[Array, "prefix_tokens channels"] | None = None,
         return_updated_state: bool = False,
     ) -> CausalConvResult:
@@ -136,9 +137,23 @@ class SeparableCausalConv(LalamoModule[SeparableCausalConvConfig]):
         if self.biases is not None:
             results = results + self.biases
 
+        if return_updated_state:
+            if length_without_padding is None:
+                length_without_padding = num_suffix_tokens
+            length_without_padding = jnp.asarray(length_without_padding, dtype=jnp.int32)
+            length_without_padding = jnp.clip(length_without_padding, 0, num_suffix_tokens)
+            updated_state = jax.lax.dynamic_slice_in_dim(
+                inputs_with_history,
+                start_index=length_without_padding,
+                slice_size=self.kernel_size - 1,
+                axis=0,
+            )
+        else:
+            updated_state = None
+
         return CausalConvResult(
             results,
-            (inputs_with_history if return_updated_state else None),
+            updated_state,
         )
 
     def export_weights(self) -> ParameterTree:
@@ -436,6 +451,7 @@ class Mamba2(TokenMixerBase[Mamba2Config, Mamba2StateLayer]):
 
         conv_output, updated_conv_state = self.conv(
             conv_inputs,
+            length_without_padding,
             state.conv_state,
             return_updated_state=return_updated_state,
         )

lalamo/modules/token_mixers/short_conv.py

@@ -116,7 +116,7 @@ class ShortConv(TokenMixerBase[ShortConvConfig, ShortConvStateLayer]):
         positional_embeddings: PositionalEmbeddings | None,
         state: ShortConvStateLayer | None = None,
         return_updated_state: bool = False,
-        length_without_padding: Int[Array, ""] | int | None = None, # noqa: ARG002
+        length_without_padding: Int[Array, ""] | int | None = None,
     ) -> TokenMixerResult[ShortConvStateLayer]:
         if positional_embeddings is not None:
             raise ValueError("Positional embeddings are not supported for ShortConv.")
@@ -124,7 +124,7 @@ class ShortConv(TokenMixerBase[ShortConvConfig, ShortConvStateLayer]):
         pre_conv_gate, post_conv_gate, x = vmap(self.in_projection)(inputs)
 
         prev_conv_state = state.conv_state if state is not None else None
-        conv_output = self.conv(x * pre_conv_gate, prev_conv_state, return_updated_state)
+        conv_output = self.conv(x * pre_conv_gate, length_without_padding, prev_conv_state, return_updated_state)
 
         (outputs,) = vmap(self.out_projection)(conv_output.outputs * post_conv_gate)
         updated_conv_state = conv_output.state

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

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

@@ -1,4 +1,4 @@
-lalamo/__init__.py,sha256=vJX2fu1UJdIdtVlBjxWtPNTw8dEge9G2PI1eQOGF-qA,1386
+lalamo/__init__.py,sha256=XBfWi6pPtdWFEQRvMxVw8KGoqWxIFq01Z2zBxqNp7BE,1386
 lalamo/commands.py,sha256=zXyyrLTHhP9wouwtpX4RUZeEF6No-_9ee-y_GWGhw7k,10972
 lalamo/common.py,sha256=WaNJx20eUX4CBF50aym9lniGAiX-SzBJzDzO5Jh6zXA,4312
 lalamo/main.py,sha256=Tez84CtMxUi1ySuRSqQElu4Zr1UWs_Gw6HX1xtCZknQ,27383
@@ -74,8 +74,8 @@ lalamo/modules/utils.py,sha256=t_TayWT6g5LtYKhJaod-u_COWaI_VbNd3eYek9Nj0lc,441
 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=d5hq1M_fk467EKm9Nyu3xu5OyrUmp04eeTyPSKh4HVI,18422
-lalamo/modules/token_mixers/short_conv.py,sha256=xLoTiN44cFuPva9gdjjC8EQqI5RYSWRd_YFxTpyOsmE,5076
+lalamo/modules/token_mixers/mamba.py,sha256=zV5CnhEbAtJ32V32a2VZGsbjZ-sohMqRbR5kW9XH1AI,19087
+lalamo/modules/token_mixers/short_conv.py,sha256=k1z9UwcJGag2NHWad7cYiAnhxULtmva9RrdhqVbir18,5085
 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
@@ -87,9 +87,9 @@ lalamo/speculator/estimator.py,sha256=S_TRwMnjWg5qt9le2AYua_Vmo6QkIT-0Si7TjCfC7x
 lalamo/speculator/inference.py,sha256=5GntUgj0HQLeLn3HIHnVX8EEO0EBzmKeP5-_U7kdFAM,3670
 lalamo/speculator/ngram.py,sha256=Fy3A-oVxZql3gE5M5ot0hKPu0772-kcEPDvD9MkldpA,5889
 lalamo/speculator/utils.py,sha256=0wZoMMIzzk0Q-3zq5H5f-JBplePNHxywndkrNtOJOyo,1697
-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,,
+lalamo-0.6.1.dist-info/licenses/LICENSE,sha256=diHRfjSEJHD1nnEeMIfMRCjR3UERf8bT3eseD6b1ayA,1072
+lalamo-0.6.1.dist-info/METADATA,sha256=eAuWPVMZl52_KExdalios28l6mOQmKgE3EcIUGUKd4k,3112
+lalamo-0.6.1.dist-info/WHEEL,sha256=qELbo2s1Yzl39ZmrAibXA2jjPLUYfnVhUNTlyF1rq0Y,92
+lalamo-0.6.1.dist-info/entry_points.txt,sha256=qli7qTfnBk5WP10rOGXXEckHMtt-atJMDWd8jN89Uks,43
+lalamo-0.6.1.dist-info/top_level.txt,sha256=VHvWL5JN5XRG36NsN_MieJ7EwRihEOrEjyDaTdFJ-aI,7
+lalamo-0.6.1.dist-info/RECORD,,

{lalamo-0.6.0.dist-info → lalamo-0.6.1.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.9.0)
+Generator: setuptools (80.10.1)
 Root-Is-Purelib: true
 Tag: py3-none-any