lalamo 0.5.9__py3-none-any.whl → 0.5.11__py3-none-any.whl

lalamo/__init__.py

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

lalamo/model_import/decoder_configs/__init__.py

@@ -6,6 +6,7 @@ from .huggingface import (
     HFGemma3Config,
     HFGemma3TextConfig,
     HFGPTOssConfig,
+    HFLFM2Config,
     HFLlamaConfig,
     HFLlambaConfig,
     HFMistralConfig,
@@ -22,6 +23,7 @@ __all__ = [
     "HFGemma2Config",
     "HFGemma3Config",
     "HFGemma3TextConfig",
+    "HFLFM2Config",
     "HFLlamaConfig",
     "HFLlambaConfig",
     "HFMistralConfig",

lalamo/model_import/decoder_configs/huggingface/__init__.py

@@ -2,6 +2,7 @@ from .common import HuggingFaceLMConfig
 from .gemma2 import HFGemma2Config
 from .gemma3 import HFGemma3Config, HFGemma3TextConfig
 from .gpt_oss import HFGPTOssConfig
+from .lfm2 import HFLFM2Config
 from .llama import HFLlamaConfig
 from .llamba import HFLlambaConfig
 from .mistral import HFMistralConfig
@@ -14,6 +15,7 @@ __all__ = [
     "HFGemma2Config",
     "HFGemma3Config",
     "HFGemma3TextConfig",
+    "HFLFM2Config",
     "HFLlamaConfig",
     "HFLlambaConfig",
     "HFMistralConfig",

lalamo/model_import/decoder_configs/huggingface/lfm2.py

@@ -0,0 +1,225 @@
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Literal
+
+from jaxtyping import DTypeLike
+
+from lalamo.modules import (
+    AttentionConfig,
+    DecoderConfig,
+    DenseMLPConfig,
+    FullPrecisionLinearConfig,
+    MLXQuantizedLinearConfig,
+    MLXQuantizedTiedEmbeddingConfig,
+    NormalizationConfig,
+    SeparableCausalConvConfig,
+    ShortConvConfig,
+    SiLU,
+    TiedEmbeddingConfig,
+    TransformerConfig,
+    TransformerLayerConfig,
+    UnscaledRoPEConfig,
+    UntiedEmbeddingConfig,
+    UpcastMode,
+)
+from lalamo.quantization import QuantizationMode
+
+from .common import HuggingFaceLMConfig
+
+
+@dataclass(frozen=True)
+class QuantizationConfig:
+    group_size: int
+    bits: int
+
+
+@dataclass(frozen=True)
+class HFLFM2Config(HuggingFaceLMConfig):
+    architectures: list[Literal["Lfm2ForCausalLM"]]
+    block_auto_adjust_ff_dim: bool
+    block_dim: int
+    block_ff_dim: int
+    block_ffn_dim_multiplier: float
+    block_mlp_init_scale: float
+    block_multiple_of: int
+    block_norm_eps: float
+    block_out_init_scale: float
+    block_use_swiglu: bool
+    block_use_xavier_init: bool
+    bos_token_id: int
+    conv_L_cache: int  # noqa: N815
+    conv_bias: bool
+    conv_dim: int
+    conv_dim_out: int
+    conv_use_xavier_init: bool
+    eos_token_id: int
+    hidden_size: int
+    initializer_range: float
+    max_position_embeddings: int
+    model_type: Literal["lfm2"]
+    norm_eps: float
+    num_attention_heads: int
+    num_heads: int
+    num_hidden_layers: int
+    num_key_value_heads: int
+    pad_token_id: int
+    rope_theta: float
+    torch_dtype: Literal["bfloat16"]
+    transformers_version: str
+    use_cache: bool
+    use_pos_enc: bool
+    vocab_size: int
+
+    intermediate_size: int | None = None
+    layer_types: list[Literal["conv", "full_attention"]] | None = None
+    full_attn_idxs: list[int] | None = None
+    tie_embedding: bool = True
+    theta: float | None = None
+
+    quantization: QuantizationConfig | None = None
+    quantization_config: QuantizationConfig | None = None
+
+    def to_decoder_config(
+        self,
+        context_length: int | None,
+        activation_precision: DTypeLike,
+        accumulation_precision: DTypeLike,
+        metadata_dict: Mapping[str, str],  # noqa: ARG002
+    ) -> DecoderConfig:
+        assert self.num_attention_heads == self.num_heads
+
+        if self.quantization_config is not None:
+            assert self.tie_embedding
+
+            embedding_config = MLXQuantizedTiedEmbeddingConfig(
+                input_scale=None,
+                logit_soft_cap=None,
+                group_size=self.quantization_config.group_size,
+                embedding_quantization_mode=QuantizationMode.from_num_bits(self.quantization_config.bits),
+                activation_quantization_mode=None,
+                activation_precision=activation_precision,
+            )
+        elif self.tie_embedding:
+            embedding_config = TiedEmbeddingConfig(
+                input_scale=None,
+                logit_soft_cap=None,
+                precision=activation_precision,
+            )
+        else:
+            embedding_config = UntiedEmbeddingConfig(
+                input_scale=None,
+                logit_soft_cap=None,
+                precision=activation_precision,
+            )
+
+        rope_config = UnscaledRoPEConfig(
+            precision=activation_precision,
+            base=self.rope_theta,
+            max_sequence_length=context_length or self.max_position_embeddings,
+        )
+
+        if self.quantization_config is None:
+            linear_config = FullPrecisionLinearConfig(activation_precision)
+        else:
+            linear_config = MLXQuantizedLinearConfig(
+                group_size=self.quantization_config.group_size,
+                weight_quantization_mode=QuantizationMode.from_num_bits(self.quantization_config.bits),
+                activation_quantization_mode=None,
+                activation_precision=activation_precision,
+            )
+
+        block_norm_config = NormalizationConfig(
+            scale_precision=activation_precision,
+            accumulation_precision=accumulation_precision,
+            epsilon=self.block_norm_eps,
+            scale_offset=None,
+            upcast_mode=UpcastMode.ONLY_NORMALIZATION,
+            subtract_mean=False,
+        )
+
+        attention_config = AttentionConfig(
+            qkv_projection_config=linear_config,
+            out_projection_config=linear_config,
+            query_norm_config=block_norm_config,
+            key_norm_config=block_norm_config,
+            num_heads=self.num_attention_heads,
+            num_groups=self.num_key_value_heads,
+            head_dim=self.hidden_size // self.num_heads,
+            is_causal=True,
+            scale=None,
+            sliding_window_size=None,
+            logit_soft_cap=None,
+            has_sinks=False,
+            has_qkv_biases=False,
+            has_out_biases=False,
+        )
+
+        short_conv_config = ShortConvConfig(
+            in_projection_config=linear_config,
+            conv_config=SeparableCausalConvConfig(activation_precision, has_biases=self.conv_bias),
+            out_projection_config=linear_config,
+            kernel_size=self.conv_L_cache,
+        )
+
+        mlp_config = DenseMLPConfig(
+            linear_config=linear_config,
+            activation=SiLU(),
+            has_up_biases=False,
+            has_down_biases=False,
+            up_clipping=None,
+            gate_clipping=None,
+        )
+
+        if self.layer_types is not None:
+            layer_types = self.layer_types
+        elif self.full_attn_idxs is not None:
+            layer_types = [
+                "full_attention" if i in self.full_attn_idxs else "conv" for i in range(self.num_hidden_layers)
+            ]
+        else:
+            raise RuntimeError("Either layer_types or full_attn_idxs must be present.")
+
+        layer_configs = [
+            TransformerLayerConfig(
+                pre_mixer_norm_config=block_norm_config,
+                mixer_config={"conv": short_conv_config, "full_attention": attention_config}[layer_type],
+                post_mixer_norm_config=None,
+                pre_mlp_norm_config=block_norm_config,
+                mlp_config=mlp_config,
+                post_mlp_norm_config=None,
+            )
+            for layer_type in layer_types
+        ]
+
+        output_norm_config = NormalizationConfig(
+            scale_precision=activation_precision,
+            accumulation_precision=accumulation_precision,
+            epsilon=self.norm_eps,
+            scale_offset=None,
+            upcast_mode=UpcastMode.ONLY_NORMALIZATION,
+            subtract_mean=False,
+        )
+
+        if self.intermediate_size is not None:
+            hidden_dim = self.intermediate_size
+        else:
+            hidden_dim_adjusted = self.block_ff_dim * self.block_ffn_dim_multiplier * (2 / 3)
+            hidden_dim = int(
+                (hidden_dim_adjusted + self.block_multiple_of - 1) // self.block_multiple_of * self.block_multiple_of,
+            )
+
+        transformer_config = TransformerConfig(
+            global_rope_config=rope_config,
+            local_rope_config=None,
+            layer_configs=tuple(layer_configs),
+            output_norm_config=output_norm_config,
+            model_dim=self.hidden_size,
+            hidden_dim=hidden_dim,
+            context_length=context_length or self.max_position_embeddings,
+        )
+
+        return DecoderConfig(
+            embedding_config=embedding_config,
+            transformer_config=transformer_config,
+            vocab_size=self.vocab_size,
+        )

lalamo/model_import/loaders/huggingface.py

@@ -8,17 +8,22 @@ from jaxtyping import Array, DTypeLike
 from lalamo.common import ParameterPath
 from lalamo.modules import (
     Attention,
+    AttentionConfig,
     Decoder,
     DenseMLP,
     FullPrecisionLinear,
     GroupQuantizedLinear,
     LinearBase,
     Mamba2,
+    Mamba2Config,
     MLXQuantizedLinear,
     MLXQuantizedTiedEmbedding,
+    MLXQuantizedTiedEmbeddingConfig,
     MLXSemiQuantizedUntiedEmbedding,
     Normalization,
     SeparableCausalConv,
+    ShortConv,
+    ShortConvConfig,
     TiedEmbedding,
     TransformerLayer,
     UntiedEmbedding,
@@ -345,21 +350,42 @@ def load_attention(
     weights_dict: Mapping[str, Array],
     path: ParameterPath,
 ) -> Attention:
+    if (path / "o_proj.weight") in weights_dict:
+        o_proj_name = "o_proj"
+    elif (path / "out_proj.weight") in weights_dict:
+        o_proj_name = "out_proj"
+    else:
+        raise NotImplementedError("Can't determine attention output projection name")
+
     qkv_projection = load_linear(
         module.qkv_projection,
         weights_dict,
         path,
         sublayers_to_fuse=["q_proj", "k_proj", "v_proj"],
     )
-    out_projection = load_linear(module.out_projection, weights_dict, path / "o_proj")
+    out_projection = load_linear(module.out_projection, weights_dict, path / o_proj_name)
 
     if module.query_norm is not None:
-        query_norm = load_rmsnorm(module.query_norm, weights_dict, path / "q_norm")
+        if (path / "q_norm.weight") in weights_dict:
+            q_norm_name = "q_norm"
+        elif (path / "q_layernorm.weight") in weights_dict:
+            q_norm_name = "q_layernorm"
+        else:
+            raise NotImplementedError("Can't determine attention query projection parameter name")
+
+        query_norm = load_rmsnorm(module.query_norm, weights_dict, path / q_norm_name)
     else:
         query_norm = None
 
     if module.key_norm is not None:
-        key_norm = load_rmsnorm(module.key_norm, weights_dict, path / "k_norm")
+        if (path / "k_norm.weight") in weights_dict:
+            k_norm_name = "k_norm"
+        elif (path / "k_layernorm.weight") in weights_dict:
+            k_norm_name = "k_layernorm"
+        else:
+            raise NotImplementedError("Can't determine attention key projection parameter name")
+
+        key_norm = load_rmsnorm(module.key_norm, weights_dict, path / k_norm_name)
     else:
         key_norm = None
 
@@ -382,19 +408,24 @@ def load_attention(
     )
 
 
-def _load_mamba_conv(
+def _load_conv(
     conv_module: SeparableCausalConv,
     weights_dict: Mapping[str, Array],
     path: ParameterPath,
+    permute_conv: bool,
 ) -> SeparableCausalConv:
     weight_path = path / "conv1d" / "weight"
     if weight_path not in weights_dict:
         weight_path = path / "conv_weight"
+    if weight_path not in weights_dict:
+        weight_path = path / "conv.weight"
     if weight_path not in weights_dict:
         weight_path = None
 
     if weight_path is not None:
         raw = weights_dict[weight_path]
+        if permute_conv:
+            raw = jnp.matrix_transpose(raw)
         conv_weight = raw.squeeze(1) if raw.ndim == 3 else raw
     else:
         conv_weight = conv_module.weights
@@ -402,6 +433,8 @@ def _load_mamba_conv(
     bias_path = path / "conv1d" / "bias"
     if bias_path not in weights_dict:
         bias_path = path / "conv_bias"
+    if bias_path not in weights_dict:
+        bias_path = path / "conv.bias"
     if bias_path not in weights_dict:
         bias_path = None
 
@@ -421,10 +454,11 @@ def load_mamba2(
     module: Mamba2,
     weights_dict: Mapping[str, Array],
     path: ParameterPath,
+    permute_conv: bool,
 ) -> Mamba2:
     in_projection = load_linear(module.in_projection, weights_dict, path / "in_proj")
     out_projection = load_linear(module.out_projection, weights_dict, path / "out_proj")
-    conv = _load_mamba_conv(module.conv, weights_dict, path)
+    conv = _load_conv(module.conv, weights_dict, path, permute_conv)
 
     skip_connection_weight_path = path / "D"
     if skip_connection_weight_path in weights_dict:
@@ -451,6 +485,23 @@ def load_mamba2(
     )
 
 
+def load_short_conv(
+    module: ShortConv,
+    weights_dict: Mapping[str, Array],
+    path: ParameterPath,
+    permute_conv: bool,
+) -> ShortConv:
+    in_projection = load_linear(module.in_projection, weights_dict, path / "in_proj")
+    out_projection = load_linear(module.out_projection, weights_dict, path / "out_proj")
+    conv = _load_conv(module.conv, weights_dict, path, permute_conv)
+
+    return load_parameters(
+        lambda m: (m.in_projection, m.out_projection, m.conv),
+        module,
+        (in_projection, out_projection, conv),
+    )
+
+
 def load_transformer_layer(
     module: TransformerLayer,
     weights_dict: Mapping[str, Array],
@@ -463,6 +514,7 @@ def load_transformer_layer(
     up_proj_key: str,
     gate_proj_key: str,
     down_proj_key: str,
+    permute_conv: bool,
 ) -> TransformerLayer:
     if module.pre_mixer_norm is not None:
         pre_attention_norm = load_rmsnorm(
@@ -477,7 +529,9 @@ def load_transformer_layer(
     if isinstance(module.mixer, Attention):
         mixer = load_attention(module.mixer, weights_dict, mixer_path / mixer_key)
     elif isinstance(module.mixer, Mamba2):
-        mixer = load_mamba2(module.mixer, weights_dict, mixer_path / mixer_key)
+        mixer = load_mamba2(module.mixer, weights_dict, mixer_path / mixer_key, permute_conv)
+    elif isinstance(module.mixer, ShortConv):
+        mixer = load_short_conv(module.mixer, weights_dict, mixer_path / mixer_key, permute_conv)
     else:
         mixer = module.mixer
 
@@ -625,11 +679,13 @@ def load_huggingface_decoder(
 
     is_llamba_full_precision = any(key.startswith("backbone.") for key in weights_dict)
     is_llamba_mlx = any(key.startswith("embedding.encoder.") for key in weights_dict)
+    is_lfm2 = any(key.startswith("model.layers.0.operator_norm.weight") for key in weights_dict)
     if is_llamba_full_precision:
         decoder_path = base_path / "backbone"
         embedding_path = decoder_path / "embedding"
         pre_mixer_norm_key = "input_layernorm"
-        mixer_key = "mixer"
+        mixer_key = {Mamba2Config: "mixer"}
+        permute_conv = False
         pre_mlp_norm_key = "post_attention_layernorm"
         mlp_key = "mlp"
         up_proj_key = "up_proj"
@@ -642,7 +698,8 @@ def load_huggingface_decoder(
         decoder_path = base_path / "model"
         embedding_path = base_path / "embedding.encoder"
         pre_mixer_norm_key = "norm"
-        mixer_key = "layer"
+        mixer_key = {Mamba2Config: "layer"}
+        permute_conv = False
         pre_mlp_norm_key = "norm"
         mlp_key = "layer"
         up_proj_key = "gate_proj"
@@ -651,11 +708,26 @@ def load_huggingface_decoder(
         alternating_layers = True
         norm_key = "norm"
         lm_head_path = base_path / "head.linear"
+    elif is_lfm2:
+        decoder_path = base_path / "model"
+        embedding_path = decoder_path / "embed_tokens"
+        pre_mixer_norm_key = "operator_norm"
+        mixer_key = {ShortConvConfig: "conv", AttentionConfig: "self_attn"}
+        permute_conv = isinstance(module.config.embedding_config, MLXQuantizedTiedEmbeddingConfig)
+        pre_mlp_norm_key = "ffn_norm"
+        mlp_key = "feed_forward"
+        up_proj_key = "w3"
+        gate_proj_key = "w1"
+        down_proj_key = "w2"
+        alternating_layers = False
+        norm_key = "embedding_norm"
+        lm_head_path = base_path / "lm_head"
     else:
         decoder_path = base_path / "model"
         embedding_path = decoder_path / "embed_tokens"
         pre_mixer_norm_key = "input_layernorm"
-        mixer_key = "self_attn"
+        mixer_key = {AttentionConfig: "self_attn"}
+        permute_conv = False
         pre_mlp_norm_key = "post_attention_layernorm"
         mlp_key = "mlp"
         up_proj_key = "up_proj"
@@ -687,13 +759,14 @@ def load_huggingface_decoder(
             weights_dict,
             decoder_path / "layers" / ((i * 2) if alternating_layers else i),
             decoder_path / "layers" / ((i * 2 + 1) if alternating_layers else i),
-            mixer_key,
+            mixer_key[type(layer.config.mixer_config)], # type: ignore
             mlp_key,
             pre_mixer_norm_key,
             pre_mlp_norm_key,
             up_proj_key,
             gate_proj_key,
             down_proj_key,
+            permute_conv,
         )
         for i, layer in enumerate(module.transformer.layers)
     )

lalamo/model_import/model_specs/__init__.py

@@ -4,6 +4,7 @@ from .essential_ai import RNJ_MODELS
 from .gemma import GEMMA_MODELS
 from .gpt_oss import GPT_OSS_MODELS
 from .huggingface import HUGGINGFACE_MODELS
+from .lfm2 import LFM2_MODELS
 from .llama import LLAMA_MODELS
 from .llamba import LLAMBA_MODELS
 from .mirai import MIRAI_CLASSIFIER_MODELS
@@ -25,6 +26,7 @@ __all__ = [
 
 
 ALL_MODEL_LISTS = [
+    LFM2_MODELS,
     LLAMA_MODELS,
     LLAMBA_MODELS,
     DEEPSEEK_MODELS,

lalamo/model_import/model_specs/common.py

@@ -56,6 +56,7 @@ class WeightsType(Enum):
                 yield MapDictValues(lambda v: cast_if_float(v, float_dtype), weights_dict), metadata_dict or {}
         else:
             import torch
+
             from lalamo.modules.torch_interop import torch_to_jax
 
             torch_weights = torch.load(filename, map_location="cpu", weights_only=True)

lalamo/model_import/model_specs/lfm2.py

@@ -0,0 +1,31 @@
+from lalamo.model_import.decoder_configs import HFLFM2Config
+from lalamo.quantization import QuantizationMode
+
+from .common import ConfigMap, FileSpec, ModelSpec
+
+__all__ = ["LFM2_MODELS"]
+
+
+def _lfm2_repo(size: str, quantization: QuantizationMode | None) -> tuple[str, str]:
+    organization = "LiquidAI" if quantization is None else "mlx-community"
+    name = f"LFM2-{size}{f'-{quantization.bits}bit' if quantization is not None else ''}"
+    return (organization, name)
+
+
+LFM2_MODELS = [
+    ModelSpec(
+        vendor="LiquidAI",
+        family="LFM2",
+        name=_lfm2_repo(size, quantization)[1],
+        size=size,
+        repo="/".join(_lfm2_repo(size, quantization)),
+        config_type=HFLFM2Config,
+        quantization=quantization,
+        configs=ConfigMap(
+            chat_template=FileSpec("chat_template.jinja"),
+        ),
+        use_cases=tuple(),
+    )
+    for size in ["350M", "700M", "1.2B", "2.6B"]
+    for quantization in [None, *([QuantizationMode.UINT4, QuantizationMode.UINT8] if size != "2.6B" else [])]
+]

lalamo/modules/__init__.py

@@ -69,6 +69,9 @@ from .token_mixers import (
     Mamba2Config,
     SeparableCausalConv,
     SeparableCausalConvConfig,
+    ShortConv,
+    ShortConvConfig,
+    ShortConvStateLayer,
     State,
     StaticKVCacheLayer,
 )
@@ -136,6 +139,9 @@ __all__ = [
     "RoutingFunction",
     "SeparableCausalConv",
     "SeparableCausalConvConfig",
+    "ShortConv",
+    "ShortConvConfig",
+    "ShortConvStateLayer",
     "SiLU",
     "SoftmaxRouting",
     "State",

lalamo/modules/token_mixers/__init__.py

@@ -3,9 +3,18 @@ from lalamo.modules.common import register_config_union
 from .attention import Attention, AttentionConfig, AttentionResult
 from .common import TokenMixerBase, TokenMixerResult
 from .mamba import Mamba2, Mamba2Config, Mamba2Result, SeparableCausalConv, SeparableCausalConvConfig
-from .state import DynamicKVCacheLayer, KVCacheLayer, Mamba2StateLayer, State, StateLayerBase, StaticKVCacheLayer
+from .short_conv import ShortConv, ShortConvConfig, ShortConvResult
+from .state import (
+    DynamicKVCacheLayer,
+    KVCacheLayer,
+    Mamba2StateLayer,
+    ShortConvStateLayer,
+    State,
+    StateLayerBase,
+    StaticKVCacheLayer,
+)
 
-TokenMixerConfig = AttentionConfig | Mamba2Config
+TokenMixerConfig = AttentionConfig | Mamba2Config | ShortConvConfig
 
 register_config_union(TokenMixerConfig)  # type: ignore (pyright bug)
 
@@ -21,6 +30,10 @@ __all__ = [
     "Mamba2StateLayer",
     "SeparableCausalConv",
     "SeparableCausalConvConfig",
+    "ShortConv",
+    "ShortConvConfig",
+    "ShortConvResult",
+    "ShortConvStateLayer",
     "State",
     "StateLayerBase",
     "StaticKVCacheLayer",

lalamo/modules/token_mixers/common.py

@@ -25,7 +25,7 @@ class TokenMixerResult[StateLayerT](NamedTuple):
 class TokenMixerConfigBase(ABC):
     @property
     @abstractmethod
-    def rope_dim(self) -> int: ...
+    def rope_dim(self) -> int | None: ...
 
     @abstractmethod
     def random_init(

lalamo/modules/token_mixers/mamba.py

@@ -184,8 +184,8 @@ class Mamba2Config(TokenMixerConfigBase):
         return self.num_heads * self.head_dim
 
     @property
-    def rope_dim(self) -> int:
-        return self.head_dim
+    def rope_dim(self) -> None:
+        return None
 
     def random_init(
         self,

lalamo/modules/token_mixers/short_conv.py

@@ -0,0 +1,168 @@
+from collections.abc import Mapping
+from dataclasses import dataclass, replace
+from typing import Self
+
+import equinox as eqx
+from jax import vmap
+from jaxtyping import Array, DTypeLike, Float, Int, PRNGKeyArray
+
+from lalamo.common import ParameterTree
+from lalamo.modules.common import PositionalEmbeddingSelector
+from lalamo.modules.linear import LinearBase, LinearConfig
+from lalamo.modules.rope import PositionalEmbeddings
+
+from .common import TokenMixerBase, TokenMixerConfigBase, TokenMixerResult
+from .mamba import SeparableCausalConv, SeparableCausalConvConfig
+from .state import ShortConvStateLayer
+
+__all__ = [
+    "ShortConv",
+    "ShortConvConfig",
+    "ShortConvResult",
+]
+
+
+ShortConvResult = TokenMixerResult[ShortConvStateLayer]
+
+
+@dataclass(frozen=True)
+class ShortConvConfig(TokenMixerConfigBase):
+    in_projection_config: LinearConfig
+    conv_config: SeparableCausalConvConfig
+    out_projection_config: LinearConfig
+
+    kernel_size: int
+
+    @property
+    def rope_dim(self) -> None:
+        return None
+
+    def random_init(
+        self,
+        model_dim: int,
+        *,
+        key: PRNGKeyArray,
+    ) -> "ShortConv":
+        in_projection = self.in_projection_config.random_init(
+            input_dim=model_dim,
+            output_dims=(model_dim,)*3,
+            has_biases=False,
+            key=key,
+        )
+
+        conv = self.conv_config.random_init(model_dim, self.kernel_size, key=key)
+
+        out_projection = self.out_projection_config.random_init(
+            input_dim=model_dim,
+            output_dims=(model_dim,),
+            has_biases=False,
+            key=key,
+        )
+
+        return ShortConv(
+            self,
+            in_projection=in_projection,
+            conv=conv,
+            out_projection=out_projection,
+        )
+
+    def empty(
+        self,
+        model_dim: int,
+    ) -> "ShortConv":
+        in_projection = self.in_projection_config.empty(
+            input_dim=model_dim,
+            output_dims=(model_dim,)*3,
+            has_biases=False,
+        )
+
+        conv = self.conv_config.empty(model_dim, self.kernel_size)
+
+        out_projection = self.out_projection_config.empty(
+            input_dim=model_dim,
+            output_dims=(model_dim,),
+            has_biases=False,
+        )
+
+        return ShortConv(
+            self,
+            in_projection=in_projection,
+            conv=conv,
+            out_projection=out_projection,
+        )
+
+
+class ShortConv(TokenMixerBase[ShortConvConfig, ShortConvStateLayer]):
+    in_projection: LinearBase
+    conv: SeparableCausalConv
+    out_projection: LinearBase
+
+    @property
+    def activation_precision(self) -> DTypeLike:
+        return self.in_projection.activation_precision
+
+    @property
+    def model_dim(self) -> int:
+        return self.in_projection.input_dim
+
+    @property
+    def positional_embedding_selector(self) -> PositionalEmbeddingSelector:
+        return PositionalEmbeddingSelector.NONE
+
+    @eqx.filter_jit
+    def __call__(
+        self,
+        inputs: Float[Array, "suffix_tokens channels"],
+        positional_embeddings: PositionalEmbeddings | None,
+        state: ShortConvStateLayer | None = None,
+        return_updated_state: bool = False,
+        length_without_padding: Int[Array, ""] | int | None = None, # noqa: ARG002
+    ) -> TokenMixerResult[ShortConvStateLayer]:
+        if positional_embeddings is not None:
+            raise ValueError("Positional embeddings are not supported for ShortConv.")
+
+        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)
+
+        (outputs,) = vmap(self.out_projection)(conv_output.outputs * post_conv_gate)
+        updated_conv_state = conv_output.state
+
+        if return_updated_state:
+            assert updated_conv_state is not None
+            updated_state = ShortConvStateLayer(updated_conv_state)
+        else:
+            updated_state = None
+
+        return TokenMixerResult(outputs, updated_state)
+
+    def init_static_state(self, capacity: int) -> ShortConvStateLayer: # noqa: ARG002
+        return ShortConvStateLayer.init(
+            self.config.kernel_size,
+            self.in_projection.input_dim,
+            self.activation_precision,
+        )
+
+    def export_weights(self) -> ParameterTree:
+        return {
+            "in_projection": self.in_projection.export_weights(),
+            "conv": self.conv.export_weights(),
+            "out_projection": self.out_projection.export_weights(),
+        }
+
+    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"]),
+        )

lalamo/modules/token_mixers/state/__init__.py

@@ -1,11 +1,13 @@
 from .common import State, StateLayerBase
 from .kv_cache import DynamicKVCacheLayer, KVCacheLayer, StaticKVCacheLayer
 from .mamba_state import Mamba2StateLayer
+from .short_conv_state import ShortConvStateLayer
 
 __all__ = [
     "DynamicKVCacheLayer",
     "KVCacheLayer",
     "Mamba2StateLayer",
+    "ShortConvStateLayer",
     "State",
     "StateLayerBase",
     "StaticKVCacheLayer",

lalamo/modules/token_mixers/state/short_conv_state.py

@@ -0,0 +1,33 @@
+from typing import Self
+
+import jax.numpy as jnp
+from jaxtyping import Array, DTypeLike, Float
+
+from lalamo.common import ParameterTree
+
+from .common import StateLayerBase
+
+__all__ = ["ShortConvStateLayer"]
+
+
+class ShortConvStateLayer(StateLayerBase):
+    conv_state: Float[Array, "*batch tokens conv_channels"]
+
+    def __post_init__(self) -> None:
+        if self.conv_state.ndim not in (2, 3):
+            raise ValueError(
+                f"Conv state must have 2 or 3 dimensions: [batch], tokens, conv_channels,"
+                f" got shape {self.conv_state.shape}",
+            )
+
+    @classmethod
+    def init(
+        cls,
+        kernel_size: int,
+        model_dim: int,
+        dtype: DTypeLike,
+    ) -> Self:
+        return cls(conv_state=jnp.zeros((kernel_size - 1, model_dim), dtype=dtype))
+
+    def export(self) -> ParameterTree:
+        return dict(conv_state=self.conv_state)

lalamo/modules/transformer.py

@@ -65,17 +65,23 @@ class TransformerConfig:
     context_length: int
 
     def random_init(self, *, key: PRNGKeyArray) -> "Transformer":
-        first_layer_config, *_ = self.layer_configs
+        rope_dims = (layer.rope_dim for layer in self.layer_configs if layer.rope_dim is not None)
+        rope_dim = next(rope_dims, None)
+        assert all(d == rope_dim for d in rope_dims)
 
         if self.global_rope_config:
+            assert rope_dim is not None
+
             global_rope = self.global_rope_config.init(
-                head_dim=first_layer_config.rope_dim,
+                head_dim=rope_dim,
                 num_timesteps=self.context_length,
             )
         else:
             global_rope = None
 
         if self.local_rope_config:
+            assert rope_dim is not None
+
             max_sliding_window_size = max(
                 layer_config.mixer_config.sliding_window_size or 0
                 for layer_config in self.layer_configs
@@ -83,7 +89,7 @@ class TransformerConfig:
             )
 
             local_rope = self.local_rope_config.init(
-                head_dim=first_layer_config.rope_dim,
+                head_dim=rope_dim,
                 num_timesteps=max(max_sliding_window_size, self.context_length),
             )
         else:
@@ -109,19 +115,25 @@ class TransformerConfig:
         )
 
     def empty(self) -> "Transformer":
-        first_layer_config, *_ = self.layer_configs
+        rope_dims = (layer.rope_dim for layer in self.layer_configs if layer.rope_dim is not None)
+        rope_dim = next(rope_dims, None)
+        assert all(d == rope_dim for d in rope_dims)
 
         if self.global_rope_config:
+            assert rope_dim is not None
+
             global_rope = self.global_rope_config.init(
-                head_dim=first_layer_config.rope_dim,
+                head_dim=rope_dim,
                 num_timesteps=self.context_length,
             )
         else:
             global_rope = None
 
         if self.local_rope_config:
+            assert rope_dim is not None
+
             local_rope = self.local_rope_config.init(
-                head_dim=first_layer_config.rope_dim,
+                head_dim=rope_dim,
                 num_timesteps=self.context_length,
             )
         else:

lalamo/modules/transformer_layer.py

@@ -89,7 +89,7 @@ class TransformerLayerConfig:
     post_mlp_norm_config: NormalizationConfig | None
 
     @property
-    def rope_dim(self) -> int:
+    def rope_dim(self) -> int | None:
         return self.mixer_config.rope_dim
 
     def random_init(

{lalamo-0.5.9.dist-info → lalamo-0.5.11.dist-info}/METADATA

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

@@ -1,4 +1,4 @@
-lalamo/__init__.py,sha256=ANgYnkcN0qtWyEPNfJb_rcAmghdwvBrHUKE2WNN0zn4,814
+lalamo/__init__.py,sha256=YYirmrwul4QLFCnBJpYyq8Os5sFQNn-9OxTjd0CbSd8,815
 lalamo/common.py,sha256=5NUFD26yQgOnEEk3LaQnce8n-VwJxILkEpFesHZhtQU,3820
 lalamo/main.py,sha256=GgUT7lT48-XQuAEH7qzsDKG8Lx9iBf-sYBIRhZL9q7E,23978
 lalamo/message_processor.py,sha256=bSUAQg7CemLTnBV4LtPxJBicAalruDCA-JXjkTYPZ8U,5797
@@ -14,14 +14,15 @@ lalamo/model_import/__init__.py,sha256=Z8pS9rbKKx1QgUy7KZtHxiNWlZhII3mdovT9d37vA
 lalamo/model_import/common.py,sha256=wvyGD-iLut_Pm3HjDMI05upqdtCW3HWeoeB0YmiFeqk,12419
 lalamo/model_import/huggingface_generation_config.py,sha256=mot6VQ6ezCtEhN6VjhnvaU-nR5P5T2BuBUgpFNnWJxU,1495
 lalamo/model_import/huggingface_tokenizer_config.py,sha256=xvwdmio7b9nhn2H3uMBVligiYj58JaCFCvHY3-8dBvM,2502
-lalamo/model_import/decoder_configs/__init__.py,sha256=1ZqMcEHvCJjMIZ9iNyY31XMXOaFxB-NbqIU01BtmcEk,641
+lalamo/model_import/decoder_configs/__init__.py,sha256=YvlSsJqNEQPCNKcUzCw0MLjt8H3vcfjc4sz1OK7qdIQ,679
 lalamo/model_import/decoder_configs/common.py,sha256=L8PCgF5fIt3RqPlmLiJpBzDguKk9iTjk4XSItxwVG4c,3260
 lalamo/model_import/decoder_configs/executorch.py,sha256=fTEG_j-7d8riR3Fu_H5tHDjOTrWevfyw7QbWF1mUdOQ,5924
-lalamo/model_import/decoder_configs/huggingface/__init__.py,sha256=3H7GPTFNNahEvI8D1SGg2mGBgPhsIdZ213MglwbGDlE,645
+lalamo/model_import/decoder_configs/huggingface/__init__.py,sha256=AboZJgZxOuIigPShskj-FqBkBqwlJZoKHP0RDqx-MyY,696
 lalamo/model_import/decoder_configs/huggingface/common.py,sha256=YYIDEQy8x7lqL2qtxUHrNqfjZEiizBZ_26sTqOzjRtQ,3792
 lalamo/model_import/decoder_configs/huggingface/gemma2.py,sha256=g8LH_GlSNyL04WWi596zI0rWsD3ahnfNjDk-9zZNcDE,4759
 lalamo/model_import/decoder_configs/huggingface/gemma3.py,sha256=aSZ0TtpgDYA10rHi8eD0C_Jsn48siM_HXqfZ4O7nh94,8372
 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/llama.py,sha256=UPeQiz2Dix8YaZYRxn9z44OZJ6c4xBQmcUZcM0Ymvh4,6934
 lalamo/model_import/decoder_configs/huggingface/llamba.py,sha256=ANB-vQK8U-zVFubZSTDXXt2S70T5SVOGzf7eOVvPzIQ,5773
 lalamo/model_import/decoder_configs/huggingface/mistral.py,sha256=MDGC0ivzJuUpOC11n8vFdcVzqccUyaRw_hkL74mVlAg,4599
@@ -31,15 +32,16 @@ lalamo/model_import/decoder_configs/huggingface/qwen3.py,sha256=lySVO-TvusAYUjDn
 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=QURyxD3C4Nzwa8k9iHVx32hQHV-aMWjb29W5_U99-WA,29834
+lalamo/model_import/loaders/huggingface.py,sha256=_vFAgz6ZR8386VIuhWAY88nJ8qoTDawCREBsRHQQmGo,32753
 lalamo/model_import/loaders/utils.py,sha256=eiX3WKFRrAfBY-dugodscNInl5o5w3KmVcgma4atpGY,2456
-lalamo/model_import/model_specs/__init__.py,sha256=8RxLEZUxpsBtTwrTUqGIwhQ-8QzOxUdx-EL__cbcTjg,1228
-lalamo/model_import/model_specs/common.py,sha256=RVPlNWHG_5OvU1W3YcOpqYz59Dh8plDmd7z1xNrqmaY,6585
+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/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=irWgylL-pc7y3Gn5DK3fjKoCT9kJWH3B7mTa-1Gmxqc,1306
 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/llama.py,sha256=Ml-xvRGlXBT9NJhmEpwgNo6C84oBSMYgA1_PrCYGcAw,990
 lalamo/model_import/model_specs/llamba.py,sha256=Ic3sWTv34FLJ4fG6OR_Mc5goGJQR6fa5b2WbVXbn9FA,1471
 lalamo/model_import/model_specs/mirai.py,sha256=eifYVV5-fABiLH6rr82_DiVFtDyqpW0vbvXCYsQQzto,617
@@ -52,7 +54,7 @@ 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/modules/__init__.py,sha256=xWJ4OPAF4gKd0evYwXIK5kTnbH6nI55oLAePcoDDHQ0,3730
+lalamo/modules/__init__.py,sha256=dFCicpcx-XV9sVTMR7x4TVF2tAGpzFi_sCTPAyawoJo,3858
 lalamo/modules/activations.py,sha256=U3qTQtZawPAUcoqbkIJnmTYcaNiQuSPMLcBeJ398GhI,1022
 lalamo/modules/classifier.py,sha256=_jtJ3INEq1dJP5HpUmcDk9YYzpRYlQ04zvFGaWBV6Lg,12101
 lalamo/modules/common.py,sha256=dqDEOi-C3H4U9iWUisU32RA-wRDCGuaUNGbObRBhyQM,3315
@@ -64,26 +66,28 @@ lalamo/modules/mlx_interop.py,sha256=FdfU_1iES-HQ9r4K0SkYwJTyvE0f-_T5ursNCjPLZKY
 lalamo/modules/normalization.py,sha256=cBdOq6OmJssunVeEwFRJD0BDhgFAN7J8gOKwzIUAY8I,3005
 lalamo/modules/rope.py,sha256=rCik7vBNqRXYg3LGbmc1mezPRNbIYMg5cydTFpQy-eU,10157
 lalamo/modules/torch_interop.py,sha256=-mujd1zI4ec2w92Hd50RtDa0K3jl6ZSnPxc5r3Fp9nU,916
-lalamo/modules/transformer.py,sha256=67-WZX2eE314abiQOhRNSooTHeJh4q9mlZQIxQ-oASU,10162
-lalamo/modules/transformer_layer.py,sha256=CfkYIn8a3pR4PPsI9hmAXpyiTbjXo-Gzl2OU9lAQlkI,12724
+lalamo/modules/transformer.py,sha256=4olEO8Eh7U6RwSnaECn39ooPuTKUZp_6QmvO6vdirrQ,10532
+lalamo/modules/transformer_layer.py,sha256=ZYmGR2Ej328l7K-YpV4eEiBk8SzLsw1RiuSiUP94UpY,12731
 lalamo/modules/utils.py,sha256=t_TayWT6g5LtYKhJaod-u_COWaI_VbNd3eYek9Nj0lc,441
-lalamo/modules/token_mixers/__init__.py,sha256=_t4T25C4WBVJQ1SqkQPGrrAc7bPKhDO3K2btgefVh5s,909
+lalamo/modules/token_mixers/__init__.py,sha256=z6x8cNjis6xIi_2llIoByKqMF2W4xJ05rDnxitHQ3jU,1139
 lalamo/modules/token_mixers/attention.py,sha256=gkGMFah2OHB_tyJpkshM1KhMnzG6U7Xt273MkBvDk58,16584
-lalamo/modules/token_mixers/common.py,sha256=-ej1pIrrp845ztavJ3oh82U3HEsV_rEHxMZOEDp7iK8,1979
-lalamo/modules/token_mixers/mamba.py,sha256=MIIMZAlVVE4YwyT0PsxA0OWXa13ondoJchRxQbhq678,18797
-lalamo/modules/token_mixers/state/__init__.py,sha256=iQaX7njz3XtwGugI5_PUOIp1wdCzd5h08UkgF6yW3zo,307
+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/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
 lalamo/modules/token_mixers/state/mamba_state.py,sha256=LHzJvNE6MkB7nrsZSNto6pxbnMJCl--JOoe9Fkcc9Mg,1642
+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/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.9.dist-info/licenses/LICENSE,sha256=diHRfjSEJHD1nnEeMIfMRCjR3UERf8bT3eseD6b1ayA,1072
-lalamo-0.5.9.dist-info/METADATA,sha256=573oeEuYV14_hFpPmW2CNVZWciVS4_V85597oKOvjpo,3146
-lalamo-0.5.9.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-lalamo-0.5.9.dist-info/entry_points.txt,sha256=qli7qTfnBk5WP10rOGXXEckHMtt-atJMDWd8jN89Uks,43
-lalamo-0.5.9.dist-info/top_level.txt,sha256=VHvWL5JN5XRG36NsN_MieJ7EwRihEOrEjyDaTdFJ-aI,7
-lalamo-0.5.9.dist-info/RECORD,,
+lalamo-0.5.11.dist-info/licenses/LICENSE,sha256=diHRfjSEJHD1nnEeMIfMRCjR3UERf8bT3eseD6b1ayA,1072
+lalamo-0.5.11.dist-info/METADATA,sha256=0uICfw4NXxiDgA8JIPaYiUnOlFcGxYe7NKbNP2ZejHE,3147
+lalamo-0.5.11.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+lalamo-0.5.11.dist-info/entry_points.txt,sha256=qli7qTfnBk5WP10rOGXXEckHMtt-atJMDWd8jN89Uks,43
+lalamo-0.5.11.dist-info/top_level.txt,sha256=VHvWL5JN5XRG36NsN_MieJ7EwRihEOrEjyDaTdFJ-aI,7
+lalamo-0.5.11.dist-info/RECORD,,