lalamo 0.6.0__tar.gz → 0.6.2__tar.gz

{lalamo-0.6.0 → lalamo-0.6.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lalamo
-Version: 0.6.0
+Version: 0.6.2
 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 → lalamo-0.6.2}/lalamo/__init__.py

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

{lalamo-0.6.0 → lalamo-0.6.2}/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-0.6.0 → lalamo-0.6.2}/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 → lalamo-0.6.2}/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].cast("f"), # noop cast to make typechecker happy
+            memoryview(self.ngram_values)[idx_start:idx_end].cast("c").cast("f"), # noop cast to make typechecker happy
             memoryview(self.ngram_counts)[seq_hash : (seq_hash + 1)],
         )
 

{lalamo-0.6.0 → lalamo-0.6.2}/lalamo.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lalamo
-Version: 0.6.0
+Version: 0.6.2
 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 → lalamo-0.6.2}/pyproject.toml

@@ -50,6 +50,7 @@ dev = [
     "ruff==0.14.0",
     "ty>=0.0.4",
     "pytest-xdist>=3.8.0",
+    "hypothesis>=6.150.0",
 ]
 
 [project.scripts]