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

lalamo/model_import/loaders/huggingface.py

@@ -289,7 +289,7 @@ def load_moe(module: MixtureOfExperts, weights_dict: Mapping[str, Array], path:
         combined_up_gate_b = jnp.concatenate([up_b + 1.0, gate_b], axis=-1)
 
         up_projection = load_parameters(
-            lambda m: (m.weights, m.biases),  # type: ignore
+            lambda m: (m.weights, m.biases),
             module.experts.up_projection,
             (combined_up_gate_w, combined_up_gate_b),
         )
@@ -309,7 +309,7 @@ def load_moe(module: MixtureOfExperts, weights_dict: Mapping[str, Array], path:
             down_b = jnp.broadcast_to(down_b, (*down_w.shape[:-1], down_b.shape[0]))
 
         down_projection = load_parameters(
-            lambda m: (m.weights, m.biases),  # type: ignore
+            lambda m: (m.weights, m.biases),
             module.experts.down_projection,
             (down_w, down_b),
         )
@@ -807,7 +807,7 @@ 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[type(layer.config.mixer_config)],  # type: ignore
+            mixer_key[type(layer.config.mixer_config)],
             mlp_key,
             pre_mixer_norm_key,
             pre_mlp_norm_key,

lalamo/model_import/model_specs/common.py

@@ -7,13 +7,14 @@ from contextlib import contextmanager
 from dataclasses import dataclass, field
 from enum import Enum, StrEnum
 from pathlib import Path
-from typing import ClassVar, cast, get_args, get_origin
+from typing import Any, ClassVar, cast, get_args, get_origin
 
 import cattrs
 import jax.numpy as jnp
 from jaxtyping import Array, DTypeLike
 
 from lalamo.model_import.decoder_configs import ForeignConfig
+from lalamo.models.language_model import GenerationConfig
 from lalamo.quantization import QuantizationMode
 from lalamo.safetensors import safe_read
 from lalamo.utils import MapDictValues
@@ -86,7 +87,7 @@ class ConfigMap:
     model_config: FileSpec = field(default=FileSpec("config.json"))
     tokenizer: FileSpec = field(default=FileSpec("tokenizer.json"))
     tokenizer_config: FileSpec = field(default=FileSpec("tokenizer_config.json"))
-    generation_config: FileSpec | None = field(default=FileSpec("generation_config.json"))
+    generation_config: FileSpec | GenerationConfig | None = field(default=FileSpec("generation_config.json"))
     chat_template: FileSpec | JSONFieldSpec | str | None = None
 
 
@@ -123,6 +124,7 @@ def _structure_chat_template(value: object, _type: object) -> FileSpec | JSONFie
     if isinstance(value, str):
         return value
     if isinstance(value, dict):
+        value = cast("dict[Any, Any]", value) # ty bug??? Why is just `dict` != `dict[Any, Any]`?
         if "file_spec" in value and "field_name" in value:
             return JSONFieldSpec(
                 file_spec=FileSpec(**value["file_spec"]),

lalamo/model_import/model_specs/lfm2.py

@@ -1,4 +1,7 @@
+from itertools import chain, product
+
 from lalamo.model_import.decoder_configs import HFLFM2Config
+from lalamo.models.language_model import GenerationConfig
 from lalamo.quantization import QuantizationMode
 
 from .common import ConfigMap, FileSpec, ModelSpec
@@ -6,26 +9,55 @@ 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)
+def _lfm_repo(family: str, size: str, variant: str | None, quantization: QuantizationMode | None) -> tuple[str, str]:
+    return (
+        "LiquidAI" if quantization is None else "mlx-community",
+        f"{family}-{size}"
+        f"{f'-{variant}' if variant is not None else ''}"
+        f"{f'-{quantization.bits}bit' if quantization is not None else ''}",
+    )
 
 
-LFM2_MODELS = [
+_LFM20_MODELS = [
     ModelSpec(
         vendor="LiquidAI",
         family="LFM2",
-        name=_lfm2_repo(size, quantization)[1],
+        name=_lfm_repo("LFM2", size, variant, quantization)[1],
         size=size,
-        repo="/".join(_lfm2_repo(size, quantization)),
+        repo="/".join(_lfm_repo("LFM2", size, variant, quantization)),
         config_type=HFLFM2Config,
         quantization=quantization,
         configs=ConfigMap(
+            generation_config=GenerationConfig(temperature=0.3, min_p=0.15),  # , repetition_penalty=1.05
             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 [])]
+    for size, variant, quantization in chain(
+        product(["350M", "700M", "1.2B"], [None], [None, QuantizationMode.UINT4, QuantizationMode.UINT8]),
+        product(["2.6B"], [None, "Exp"], [None]),
+        product(["2.6B"], ["Exp"], [QuantizationMode.UINT4, QuantizationMode.UINT8]),
+    )
 ]
+
+_LFM25_MODELS = [
+    ModelSpec(
+        vendor="LiquidAI",
+        family="LFM2.5",
+        name=_lfm_repo("LFM2.5", size, variant, quantization)[1],
+        size=size,
+        repo="/".join(_lfm_repo("LFM2.5", size, variant, quantization)),
+        config_type=HFLFM2Config,
+        quantization=quantization,
+        configs=ConfigMap(
+            generation_config=GenerationConfig(temperature=0.1, top_k=50, top_p=0.1),  # , repetition_penalty=1.05
+            chat_template=FileSpec("chat_template.jinja"),
+        ),
+        use_cases=tuple(),
+    )
+    for size, variant, quantization in chain(
+        product(["1.2B"], ["Instruct"], [None]),
+    )
+]
+
+LFM2_MODELS = _LFM20_MODELS + _LFM25_MODELS

lalamo/models/language_model.py

@@ -64,14 +64,15 @@ class GenerationResults(NamedTuple):
 
 @dataclass(frozen=True)
 class GenerationConfig:
-    stop_token_ids: tuple[int, ...]
-    temperature: float | None
-    top_k: int | None
-    top_p: float | None
-    banned_tokens: tuple[int, ...] | None
+    stop_token_ids: tuple[int, ...] = tuple()
+    temperature: float | None = None
+    top_k: int | None = None
+    top_p: float | None = None
+    min_p: float | None = None
+    banned_tokens: tuple[int, ...] | None = None
 
     def default_policy(self) -> SamplingPolicy:
-        return make_policy(self.temperature, self.top_k, self.top_p, self.banned_tokens)
+        return make_policy(self.temperature, self.top_k, self.top_p, self.min_p, self.banned_tokens)
 
 
 @dataclass(frozen=True)

lalamo/modules/activations.py

@@ -44,4 +44,4 @@ class Identity(ActivationBase):
 Activation = SiLU | GELU | Identity
 
 
-register_config_union(Activation)  # type: ignore (pyright bug)
+register_config_union(Activation)

lalamo/modules/classifier.py

@@ -9,7 +9,7 @@ from jax import numpy as jnp
 from jax import vmap
 from jaxtyping import Array, DTypeLike, Float, Int, PRNGKeyArray
 
-from lalamo.common import ParameterTree
+from lalamo.common import ParameterTree, require_tree
 from lalamo.modules import Activation
 from lalamo.modules.normalization import NormalizationConfig
 from lalamo.modules.transformer import (
@@ -67,7 +67,7 @@ class PredictionHeadConfig:
     def random_init(self, input_size: int, num_labels: int, key: PRNGKeyArray) -> "PredictionHead":
         dense_key, readout_key = jax.random.split(key)
         dense_layer = self.dense_config.random_init(
-            input_size, (input_size,), has_biases=self.use_dense_bias, key=dense_key
+            input_size, (input_size,), has_biases=self.use_dense_bias, key=dense_key,
         )
         norm = self.normalization_config.empty(input_size)
         readout = self.readout_config.random_init(
@@ -117,19 +117,13 @@ class PredictionHead(LalamoModule[PredictionHeadConfig]):
         )
         return result
 
-    def import_weights(
-        self,
-        weights: ParameterTree[Array],
-    ) -> Self:
+    def import_weights(self, weights: ParameterTree[Array]) -> Self:
         assert isinstance(weights, Mapping)
-        assert isinstance(weights["dense"], Mapping)
-        assert isinstance(weights["norm"], Mapping)
-        assert isinstance(weights["readout"], Mapping)
         return replace(
             self,
-            dense=self.dense.import_weights(weights["dense"]),
-            norm=self.norm.import_weights(weights["norm"]),
-            readout=self.readout.import_weights(weights["readout"]),
+            dense=self.dense.import_weights(require_tree(weights["dense"])),
+            norm=self.norm.import_weights(require_tree(weights["norm"])),
+            readout=self.readout.import_weights(require_tree(weights["readout"])),
         )
 
 
@@ -321,19 +315,12 @@ class Classifier(LalamoModule[ClassifierConfig]):
         )
         return result
 
-    def import_weights(
-        self,
-        weights: ParameterTree[Array],
-    ) -> Self:
+    def import_weights(self, weights: ParameterTree[Array]) -> Self:
         assert isinstance(weights, Mapping)
-        assert isinstance(weights["embedding"], Mapping)
-        assert isinstance(weights["embedding_norm"], Mapping)
-        assert isinstance(weights["transformer"], Mapping)
-        assert isinstance(weights["prediction_head"], Mapping)
         return replace(
             self,
-            embedding=self.embedding.import_weights(weights["embedding"]),
-            embedding_norm=self.embedding_norm.import_weights(weights["embedding_norm"]),
-            transformer=self.transformer.import_weights(weights["transformer"]),
-            prediction_head=self.prediction_head.import_weights(weights["prediction_head"]),
+            embedding=self.embedding.import_weights(require_tree(weights["embedding"])),
+            embedding_norm=self.embedding_norm.import_weights(require_tree(weights["embedding_norm"])),
+            transformer=self.transformer.import_weights(require_tree(weights["transformer"])),
+            prediction_head=self.prediction_head.import_weights(require_tree(weights["prediction_head"])),
         )

lalamo/modules/common.py

@@ -9,15 +9,18 @@ from cattrs import Converter
 from jax import numpy as jnp
 from jaxtyping import Array, DTypeLike
 
-from lalamo.common import ParameterTree
+from lalamo.common import ParameterTree, require_array, require_tree
 
 __all__ = [
     "DummyUnionMember",
     "ForwardPassMode",
     "LalamoModule",
+    "ParameterTree",
     "PositionalEmbeddingSelector",
     "config_converter",
     "register_config_union",
+    "require_array",
+    "require_tree",
 ]
 
 

lalamo/modules/decoder.py

@@ -7,7 +7,7 @@ import jax
 from jax import vmap
 from jaxtyping import Array, DTypeLike, Float, Int, PRNGKeyArray
 
-from lalamo.common import ParameterTree
+from lalamo.common import ParameterTree, require_tree
 
 from .common import ForwardPassMode, LalamoModule
 from .embedding import EmbeddingBase, EmbeddingConfig
@@ -126,7 +126,7 @@ class Decoder(LalamoModule[DecoderConfig]):
         return self.embedding.activation_precision
 
     @eqx.filter_jit
-    def __call__(  # noqa: PLR0912
+    def __call__(
         self,
         token_ids: Int[Array, "batch suffix_tokens"],
         token_positions: Int[Array, "batch suffix_tokens"],
@@ -193,16 +193,10 @@ class Decoder(LalamoModule[DecoderConfig]):
             transformer=self.transformer.export_weights(),
         )
 
-    def import_weights(
-        self,
-        weights: ParameterTree[Array],
-    ) -> Self:
+    def import_weights(self, weights: ParameterTree[Array]) -> Self:
         assert isinstance(weights, Mapping)
-        assert isinstance(weights["embedding"], Mapping)
-        assert isinstance(weights["transformer"], Mapping)
-
         return replace(
             self,
-            embedding=self.embedding.import_weights(weights["embedding"]),
-            transformer=self.transformer.import_weights(weights["transformer"]),
+            embedding=self.embedding.import_weights(require_tree(weights["embedding"])),
+            transformer=self.transformer.import_weights(require_tree(weights["transformer"])),
         )

lalamo/modules/embedding.py

@@ -9,7 +9,7 @@ import jax.numpy as jnp
 from einops import rearrange
 from jaxtyping import Array, DTypeLike, Float, Int, PRNGKeyArray
 
-from lalamo.common import ParameterTree, dummy_array
+from lalamo.common import ParameterTree, dummy_array, require_array
 from lalamo.quantization import QuantizationMode, dynamically_quantize_activations, quantize_weights
 from lalamo.utils import jax_uint4_to_packed_uint8, jax_uint8_to_unpacked_uint4
 
@@ -355,21 +355,15 @@ class QuantizedTiedEmbedding(EmbeddingBase[QuantizedTiedEmbeddingConfig]):
             "scales": self.scales,
         }
 
-    def import_weights(
-        self,
-        weights: ParameterTree[Array],
-    ) -> Self:
+    def import_weights(self, weights: ParameterTree[Array]) -> Self:
         assert isinstance(weights, Mapping)
-        assert isinstance(weights["weights"], Array)
-        stored_weights = weights["weights"]
-
+        stored_weights = require_array(weights["weights"])
         if self.config.embedding_quantization_mode == QuantizationMode.UINT4:
             stored_weights = jax_uint8_to_unpacked_uint4(stored_weights)
-
         return replace(
             self,
             weights=stored_weights.astype(self.weights.dtype),
-            scales=weights["scales"],
+            scales=require_array(weights["scales"]),
         )
 
 
@@ -472,25 +466,16 @@ class MLXQuantizedTiedEmbedding(EmbeddingBase[MLXQuantizedTiedEmbeddingConfig]):
             "biases": self.biases,
         }
 
-    def import_weights(
-        self,
-        weights: ParameterTree[Array],
-    ) -> Self:
+    def import_weights(self, weights: ParameterTree[Array]) -> Self:
         assert isinstance(weights, Mapping)
-        assert isinstance(weights["weights"], Array)
-        assert isinstance(weights["scales"], Array)
-        assert isinstance(weights["biases"], Array)
-
-        unpacked_weights = weights["weights"]
-
+        unpacked_weights = require_array(weights["weights"])
         if self.config.embedding_quantization_mode == QuantizationMode.UINT4:
-            unpacked_weights = jax_uint8_to_unpacked_uint4(weights["weights"])
-
+            unpacked_weights = jax_uint8_to_unpacked_uint4(unpacked_weights)
         return replace(
             self,
             weights=unpacked_weights.astype(self.weights.dtype),
-            scales=weights["scales"],
-            biases=weights["biases"],
+            scales=require_array(weights["scales"]),
+            biases=require_array(weights["biases"]),
         )
 
 
@@ -630,33 +615,21 @@ class MLXQuantizedUntiedEmbedding(EmbeddingBase[MLXQuantizedUntiedEmbeddingConfi
             "output_biases": self.output_biases,
         }
 
-    def import_weights(
-        self,
-        weights: ParameterTree[Array],
-    ) -> Self:
+    def import_weights(self, weights: ParameterTree[Array]) -> Self:
         assert isinstance(weights, Mapping)
-        assert isinstance(weights["input_weights"], Array)
-        assert isinstance(weights["input_scales"], Array)
-        assert isinstance(weights["input_biases"], Array)
-        assert isinstance(weights["output_weights"], Array)
-        assert isinstance(weights["output_scales"], Array)
-        assert isinstance(weights["output_biases"], Array)
-
-        unpacked_input_weights = weights["input_weights"]
-        unpacked_output_weights = weights["output_weights"]
-
+        unpacked_input_weights = require_array(weights["input_weights"])
+        unpacked_output_weights = require_array(weights["output_weights"])
         if self.config.embedding_quantization_mode == QuantizationMode.UINT4:
-            unpacked_input_weights = jax_uint8_to_unpacked_uint4(weights["input_weights"])
-            unpacked_output_weights = jax_uint8_to_unpacked_uint4(weights["output_weights"])
-
+            unpacked_input_weights = jax_uint8_to_unpacked_uint4(unpacked_input_weights)
+            unpacked_output_weights = jax_uint8_to_unpacked_uint4(unpacked_output_weights)
         return replace(
             self,
             input_weights=unpacked_input_weights.astype(self.input_weights.dtype),
-            input_scales=weights["input_scales"],
-            input_biases=weights["input_biases"],
+            input_scales=require_array(weights["input_scales"]),
+            input_biases=require_array(weights["input_biases"]),
             output_weights=unpacked_output_weights.astype(self.output_weights.dtype),
-            output_scales=weights["output_scales"],
-            output_biases=weights["output_biases"],
+            output_scales=require_array(weights["output_scales"]),
+            output_biases=require_array(weights["output_biases"]),
         )
 
 
@@ -765,27 +738,17 @@ class MLXSemiQuantizedUntiedEmbedding(EmbeddingBase[MLXSemiQuantizedUntiedEmbedd
             "output_biases": self.output_biases,
         }
 
-    def import_weights(
-        self,
-        weights: ParameterTree[Array],
-    ) -> Self:
+    def import_weights(self, weights: ParameterTree[Array]) -> Self:
         assert isinstance(weights, Mapping)
-        assert isinstance(weights["input_weights"], Array)
-        assert isinstance(weights["output_weights"], Array)
-        assert isinstance(weights["output_scales"], Array)
-        assert isinstance(weights["output_biases"], Array)
-
-        unpacked_output_weights = weights["output_weights"]
-
+        unpacked_output_weights = require_array(weights["output_weights"])
         if self.config.embedding_quantization_mode == QuantizationMode.UINT4:
-            unpacked_output_weights = jax_uint8_to_unpacked_uint4(weights["output_weights"])
-
+            unpacked_output_weights = jax_uint8_to_unpacked_uint4(unpacked_output_weights)
         return replace(
             self,
-            input_weights=weights["input_weights"],
+            input_weights=require_array(weights["input_weights"]),
             output_weights=unpacked_output_weights.astype(self.output_weights.dtype),
-            output_scales=weights["output_scales"],
-            output_biases=weights["output_biases"],
+            output_scales=require_array(weights["output_scales"]),
+            output_biases=require_array(weights["output_biases"]),
         )
 
 
@@ -799,4 +762,4 @@ EmbeddingConfig = (
 )
 
 
-register_config_union(EmbeddingConfig)  # type: ignore (pyright bug)
+register_config_union(EmbeddingConfig)

lalamo/modules/linear.py

@@ -2,7 +2,7 @@ import math
 from abc import ABC, abstractmethod
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass, replace
-from typing import Self
+from typing import Self, cast
 
 import equinox as eqx
 import jax
@@ -10,7 +10,7 @@ import jax.numpy as jnp
 from einops import rearrange
 from jaxtyping import Array, DTypeLike, Float, Int, PRNGKeyArray
 
-from lalamo.common import ParameterTree, dummy_array
+from lalamo.common import ParameterTree, dummy_array, require_array
 from lalamo.quantization import QuantizationMode, dynamically_quantize_activations, quantize_weights
 from lalamo.utils import jax_uint4_to_packed_uint8, jax_uint8_to_unpacked_uint4
 
@@ -464,7 +464,7 @@ class GroupQuantizedLinearBase[ConfigT: GroupQuantizedLinearConfig](QuantizedLin
 
         return packed
 
-    def __post_init__(self) -> None:  # noqa: PLR0912
+    def __post_init__(self) -> None:
         if self.weights.dtype != self.config.activation_precision:
             raise ValueError(
                 f"Weight dtype ({self.weights.dtype}) is not equal to specified activation precision"
@@ -572,26 +572,19 @@ class GroupQuantizedLinearBase[ConfigT: GroupQuantizedLinearConfig](QuantizedLin
             result["biases"] = self.biases
         return result
 
-    def import_weights(
-        self,
-        weights: ParameterTree[Array],
-    ) -> Self:
+    def import_weights(self, weights: ParameterTree[Array]) -> Self:
         assert isinstance(weights, Mapping)
-        assert isinstance(weights["weights"], Array)
-        assert isinstance(weights["zero_points"], Array)
-        unpacked_weights = weights["weights"]
-        unpacked_zero_points = weights["zero_points"]
-
+        unpacked_weights = require_array(weights["weights"])
+        unpacked_zero_points = require_array(weights["zero_points"])
         if self.config.weight_quantization_mode == QuantizationMode.UINT4:
-            unpacked_weights = jax_uint8_to_unpacked_uint4(weights["weights"])
-            unpacked_zero_points = jax_uint8_to_unpacked_uint4(weights["zero_points"])
-
+            unpacked_weights = jax_uint8_to_unpacked_uint4(unpacked_weights)
+            unpacked_zero_points = jax_uint8_to_unpacked_uint4(unpacked_zero_points)
         return replace(
             self,
             weights=unpacked_weights.astype(self.weights.dtype),
-            scales=weights["scales"],
+            scales=require_array(weights["scales"]),
             zero_points=unpacked_zero_points.astype(self.zero_points.dtype),
-            biases=weights["biases"] if self.has_biases else None,
+            biases=require_array(weights["biases"]) if self.has_biases else None,
         )
 
 
@@ -740,7 +733,7 @@ class MLXQuantizedLinearBase[ConfigT: MLXQuantizedLinearConfig](QuantizedLinearB
 
         return packed
 
-    def __post_init__(self) -> None:  # noqa: PLR0912
+    def __post_init__(self) -> None:
         if self.weights.dtype != self.config.activation_precision:
             raise ValueError(
                 f"Weight dtype ({self.weights.dtype}) is not equal to specified activation precision"
@@ -847,24 +840,17 @@ class MLXQuantizedLinearBase[ConfigT: MLXQuantizedLinearConfig](QuantizedLinearB
             result["biases"] = self.biases
         return result
 
-    def import_weights(
-        self,
-        weights: ParameterTree[Array],
-    ) -> Self:
+    def import_weights(self, weights: ParameterTree[Array]) -> Self:
         assert isinstance(weights, Mapping)
-        assert isinstance(weights["weights"], Array)
-
-        unpacked_weights = weights["weights"]
-
+        unpacked_weights = require_array(weights["weights"])
         if self.config.weight_quantization_mode == QuantizationMode.UINT4:
-            unpacked_weights = jax_uint8_to_unpacked_uint4(weights["weights"])
-
+            unpacked_weights = jax_uint8_to_unpacked_uint4(unpacked_weights)
         return replace(
             self,
             weights=unpacked_weights.astype(self.weights.dtype),
-            scales=weights["scales"],
-            deq_biases=weights["deq_biases"],
-            biases=weights["biases"] if self.has_biases else None,
+            scales=require_array(weights["scales"]),
+            deq_biases=require_array(weights["deq_biases"]),
+            biases=require_array(weights["biases"]) if self.has_biases else None,
         )
 
 
@@ -1113,7 +1099,7 @@ class QLoRALinear(GroupQuantizedLinearBase[QLoRALinearConfig]):
         self,
         weights: ParameterTree[Array],
     ) -> Self:
-        base = super().import_weights(weights)
+        base = cast("Self", super().import_weights(weights)) # ty bug
         assert isinstance(weights, Mapping)
         assert isinstance(weights["up_weights"], Sequence)
         return replace(
@@ -1126,4 +1112,4 @@ class QLoRALinear(GroupQuantizedLinearBase[QLoRALinearConfig]):
 LinearConfig = FullPrecisionLinearConfig | GroupQuantizedLinearConfig | MLXQuantizedLinearConfig | QLoRALinearConfig
 
 
-register_config_union(LinearConfig)  # type: ignore (pyright bug)
+register_config_union(LinearConfig)

lalamo/modules/mlp.py

@@ -12,7 +12,7 @@ from einops import rearrange
 from jax import vmap
 from jaxtyping import Array, Bool, DTypeLike, Float, Int, PRNGKeyArray
 
-from lalamo.common import ParameterTree
+from lalamo.common import ParameterTree, require_tree
 from lalamo.modules.utils import vmap_twice
 
 from .activations import Activation
@@ -242,17 +242,12 @@ class DenseMLP(MLPBase[DenseMLPConfig]):
             "down_projection": self.down_projection.export_weights(),
         }
 
-    def import_weights(
-        self,
-        weights: ParameterTree[Array],
-    ) -> Self:
+    def import_weights(self, weights: ParameterTree[Array]) -> Self:
         assert isinstance(weights, Mapping)
-        assert isinstance(weights["up_projection"], Mapping)
-        assert isinstance(weights["down_projection"], Mapping)
         return replace(
             self,
-            up_projection=self.up_projection.import_weights(weights["up_projection"]),
-            down_projection=self.down_projection.import_weights(weights["down_projection"]),
+            up_projection=self.up_projection.import_weights(require_tree(weights["up_projection"])),
+            down_projection=self.down_projection.import_weights(require_tree(weights["down_projection"])),
         )
 
 
@@ -285,7 +280,7 @@ class SoftmaxRouting(RoutingFunctionBase):
 RoutingFunction = SoftmaxRouting | DummyUnionMember
 
 
-register_config_union(RoutingFunction)  # type: ignore (pyright bug)
+register_config_union(RoutingFunction)
 
 
 @dataclass(frozen=True)
@@ -486,21 +481,16 @@ class MixtureOfExperts(MLPBase[MixtureOfExpertsConfig]):
             "experts": self.experts.export_weights(),
         }
 
-    def import_weights(
-        self,
-        weights: ParameterTree[Array],
-    ) -> Self:
+    def import_weights(self, weights: ParameterTree[Array]) -> Self:
         assert isinstance(weights, Mapping)
-        assert isinstance(weights["router"], Mapping)
-        assert isinstance(weights["experts"], Mapping)
         return replace(
             self,
-            router=self.router.import_weights(weights["router"]),
-            experts=self.experts.import_weights(weights["experts"]),
+            router=self.router.import_weights(require_tree(weights["router"])),
+            experts=self.experts.import_weights(require_tree(weights["experts"])),
         )
 
 
 MLPConfig = DenseMLPConfig | MixtureOfExpertsConfig
 
 
-register_config_union(MLPConfig)  # type: ignore (pyright bug)
+register_config_union(MLPConfig)

lalamo/modules/mlx_interop.py

@@ -1,5 +1,5 @@
 import jax.numpy as jnp
-import mlx.core as mx
+import mlx.core as mx  # type: ignore
 from jaxtyping import Array
 
 __all__ = ["jax_to_mlx", "mlx_to_jax"]

lalamo/modules/rope.py

@@ -281,4 +281,4 @@ class LinearScalingRoPEConfig(RoPEConfigBase):
 
 RoPEConfig = UnscaledRoPEConfig | LlamaRoPEConfig | YARNRoPEConfig | LinearScalingRoPEConfig
 
-register_config_union(RoPEConfig)  # type: ignore (pyright bug)
+register_config_union(RoPEConfig)

lalamo/modules/token_mixers/__init__.py

@@ -16,7 +16,7 @@ from .state import (
 
 TokenMixerConfig = AttentionConfig | Mamba2Config | ShortConvConfig
 
-register_config_union(TokenMixerConfig)  # type: ignore (pyright bug)
+register_config_union(TokenMixerConfig)
 
 __all__ = [
     "Attention",