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

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",

lalamo/modules/token_mixers/attention.py

@@ -10,7 +10,7 @@ from jax import vmap
 from jaxtyping import Array, Bool, DTypeLike, Float, Int, PRNGKeyArray
 
 from lalamo.common import dummy_array
-from lalamo.modules.common import ParameterTree, PositionalEmbeddingSelector
+from lalamo.modules.common import ParameterTree, PositionalEmbeddingSelector, require_array, require_tree
 from lalamo.modules.linear import LinearBase, LinearConfig
 from lalamo.modules.normalization import Normalization, NormalizationConfig
 from lalamo.modules.rope import PositionalEmbeddings
@@ -433,33 +433,15 @@ class Attention(TokenMixerBase[AttentionConfig, KVCacheLayer]):
             result["sinks"] = self.sinks
         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["qkv_projection"], Mapping)
-        assert isinstance(weights["out_projection"], Mapping)
-        if self.query_norm is not None:
-            assert isinstance(weights["query_norm"], Mapping)
-            query_norm = self.query_norm.import_weights(weights["query_norm"])
-        else:
-            query_norm = None
-        if self.key_norm is not None:
-            assert isinstance(weights["key_norm"], Mapping)
-            key_norm = self.key_norm.import_weights(weights["key_norm"])
-        else:
-            key_norm = None
-        if self.sinks is not None:
-            assert isinstance(weights["sinks"], Array)
-            sinks = weights["sinks"]
-        else:
-            sinks = None
         return replace(
             self,
-            qkv_projection=self.qkv_projection.import_weights(weights["qkv_projection"]),
-            out_projection=self.out_projection.import_weights(weights["out_projection"]),
-            query_norm=query_norm,
-            key_norm=key_norm,
-            sinks=sinks,
+            qkv_projection=self.qkv_projection.import_weights(require_tree(weights["qkv_projection"])),
+            out_projection=self.out_projection.import_weights(require_tree(weights["out_projection"])),
+            query_norm=self.query_norm.import_weights(require_tree(weights["query_norm"]))
+            if self.query_norm
+            else None,
+            key_norm=self.key_norm.import_weights(require_tree(weights["key_norm"])) if self.key_norm else None,
+            sinks=require_array(weights["sinks"]) if self.sinks is not None else None,
         )

lalamo/modules/token_mixers/mamba.py

@@ -10,7 +10,7 @@ from einops import einsum, rearrange
 from jax import vmap
 from jaxtyping import Array, DTypeLike, Float, Int, PRNGKeyArray
 
-from lalamo.common import ParameterTree, dummy_array
+from lalamo.common import ParameterTree, dummy_array, require_array, require_tree
 from lalamo.modules.activations import Activation
 from lalamo.modules.common import LalamoModule, PositionalEmbeddingSelector
 from lalamo.modules.linear import LinearBase, LinearConfig
@@ -149,16 +149,10 @@ class SeparableCausalConv(LalamoModule[SeparableCausalConvConfig]):
 
     def import_weights(self, weights: ParameterTree[Array]) -> "SeparableCausalConv":
         assert isinstance(weights, Mapping)
-        assert isinstance(weights["weights"], Array)
-        if self.biases is not None:
-            assert isinstance(weights["biases"], Array)
-            biases = weights["biases"]
-        else:
-            biases = None
         return replace(
             self,
-            weights=weights["weights"],
-            biases=biases,
+            weights=require_array(weights["weights"]),
+            biases=require_array(weights["biases"]) if self.biases is not None else None,
         )
 
 
@@ -532,22 +526,13 @@ class Mamba2(TokenMixerBase[Mamba2Config, Mamba2StateLayer]):
             "gate_bias": self.gate_bias,
         }
 
-    def import_weights(
-        self,
-        weights: ParameterTree[Array],
-    ) -> Self:
+    def import_weights(self, weights: ParameterTree[Array]) -> Self:
         assert isinstance(weights, Mapping)
-        assert isinstance(weights["in_projection"], Mapping)
-        assert isinstance(weights["out_projection"], Mapping)
-        assert isinstance(weights["conv"], Mapping)
-        assert isinstance(weights["skip_connection_weight"], Array)
-        assert isinstance(weights["gate_bias"], Array)
-
         return replace(
             self,
-            in_projection=self.in_projection.import_weights(weights["in_projection"]),
-            out_projection=self.out_projection.import_weights(weights["out_projection"]),
-            conv=self.conv.import_weights(weights["conv"]),
-            skip_connection_weight=weights["skip_connection_weight"],
-            gate_bias=weights["gate_bias"],
+            in_projection=self.in_projection.import_weights(require_tree(weights["in_projection"])),
+            out_projection=self.out_projection.import_weights(require_tree(weights["out_projection"])),
+            conv=self.conv.import_weights(require_tree(weights["conv"])),
+            skip_connection_weight=require_array(weights["skip_connection_weight"]),
+            gate_bias=require_array(weights["gate_bias"]),
         )

lalamo/modules/token_mixers/short_conv.py

@@ -6,7 +6,7 @@ import equinox as eqx
 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.common import PositionalEmbeddingSelector
 from lalamo.modules.linear import LinearBase, LinearConfig
 from lalamo.modules.rope import PositionalEmbeddings
@@ -151,18 +151,11 @@ class ShortConv(TokenMixerBase[ShortConvConfig, ShortConvStateLayer]):
             "out_projection": self.out_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["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"]),
+            in_projection=self.in_projection.import_weights(require_tree(weights["in_projection"])),
+            conv=self.conv.import_weights(require_tree(weights["conv"])),
+            out_projection=self.out_projection.import_weights(require_tree(weights["out_projection"])),
         )

lalamo/modules/transformer.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 lalamo.modules.token_mixers import AttentionConfig
 from lalamo.modules.utils import vmap_twice
 
@@ -182,7 +182,8 @@ class Transformer(LalamoModule[TransformerConfig]):
     ) -> TransformerResult:
         if inner_features.ndim != 3:
             raise ValueError(
-                f"inner_features must be a 3D array of size (batch_size, sequence_length, hidden_dim), got {inner_features.shape}",
+                "inner_features must be a 3D array of size (batch_size, sequence_length, hidden_dim),"
+                f" got {inner_features.shape}",
             )
         if token_positions.ndim != 2:
             raise ValueError(
@@ -251,35 +252,24 @@ class Transformer(LalamoModule[TransformerConfig]):
             result["local_rope"] = self.local_rope.export_weights()
         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["layers"], Sequence)
-        assert isinstance(weights["output_norm"], Mapping)
-
         if self.global_rope:
-            assert isinstance(weights["global_rope"], Mapping)
-            global_rope = self.global_rope.import_weights(weights["global_rope"])
+            global_rope = self.global_rope.import_weights(require_tree(weights["global_rope"]))
         else:
             global_rope = None
-
         if self.local_rope:
-            assert isinstance(weights["local_rope"], Mapping)
-            local_rope = self.local_rope.import_weights(weights["local_rope"])
+            local_rope = self.local_rope.import_weights(require_tree(weights["local_rope"]))
         else:
             local_rope = None
-
-        layers = []
-        for layer, layer_weights in zip(self.layers, weights["layers"], strict=True):
-            assert isinstance(layer_weights, Mapping)
-            layers.append(layer.import_weights(layer_weights))
-
+        layers = [
+            layer.import_weights(require_tree(lw)) for layer, lw in zip(self.layers, weights["layers"], strict=True)
+        ]
         return replace(
             self,
             global_rope=global_rope,
             layers=tuple(layers),
-            output_norm=self.output_norm.import_weights(weights["output_norm"]),
+            output_norm=self.output_norm.import_weights(require_tree(weights["output_norm"])),
             local_rope=local_rope,
         )

lalamo/modules/transformer_layer.py

@@ -9,7 +9,7 @@ import jax.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 .common import ForwardPassMode, LalamoModule, PositionalEmbeddingSelector
 from .mlp import MLPBase, MLPConfig, MLPForwardPassConfig
@@ -293,38 +293,26 @@ class TransformerLayer(LalamoModule[TransformerLayerConfig]):
             result["post_mlp_norm"] = self.post_mlp_norm.export_weights()
         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["mixer"], Mapping)
-        assert isinstance(weights["mlp"], Mapping)
-        assert isinstance(weights["pre_mlp_norm"], Mapping)
-
         if self.post_mixer_norm is not None:
-            assert isinstance(weights["post_mixer_norm"], Mapping)
-            post_mixer_norm = self.post_mixer_norm.import_weights(
-                weights["post_mixer_norm"],
-            )
+            post_mixer_norm = self.post_mixer_norm.import_weights(require_tree(weights["post_mixer_norm"]))
         else:
             post_mixer_norm = None
         if self.post_mlp_norm is not None:
-            assert isinstance(weights["post_mlp_norm"], Mapping)
-            post_mlp_norm = self.post_mlp_norm.import_weights(weights["post_mlp_norm"])
+            post_mlp_norm = self.post_mlp_norm.import_weights(require_tree(weights["post_mlp_norm"]))
         else:
             post_mlp_norm = None
         if self.pre_mixer_norm is not None:
-            assert isinstance(weights["pre_mixer_norm"], Mapping)
-            pre_mixer_norm = self.pre_mixer_norm.import_weights(weights["pre_mixer_norm"])
+            pre_mixer_norm = self.pre_mixer_norm.import_weights(require_tree(weights["pre_mixer_norm"]))
         else:
             pre_mixer_norm = None
         return replace(
             self,
             pre_mixer_norm=pre_mixer_norm,
-            mixer=self.mixer.import_weights(weights["mixer"]),
+            mixer=self.mixer.import_weights(require_tree(weights["mixer"])),
             post_mixer_norm=post_mixer_norm,
-            pre_mlp_norm=self.pre_mlp_norm.import_weights(weights["pre_mlp_norm"]),
-            mlp=self.mlp.import_weights(weights["mlp"]),
+            pre_mlp_norm=self.pre_mlp_norm.import_weights(require_tree(weights["pre_mlp_norm"])),
+            mlp=self.mlp.import_weights(require_tree(weights["mlp"])),
             post_mlp_norm=post_mlp_norm,
         )

lalamo/registry_abc.py

@@ -1,5 +1,5 @@
 from abc import ABC, ABCMeta
-from typing import Any
+from typing import Any, Self
 from weakref import WeakSet
 
 __all__ = ["RegistryABC", "RegistryMeta"]
@@ -29,7 +29,7 @@ class RegistryMeta(ABCMeta):
 
         # Detect and remember the root exactly once
         if RegistryMeta._ROOT is None and name == "RegistryABC":
-            RegistryMeta._ROOT = cls  # type: ignore[assignment]
+            RegistryMeta._ROOT = cls
             return
 
         root = RegistryMeta._ROOT
@@ -58,6 +58,6 @@ class RegistryABC(ABC, metaclass=RegistryMeta):
     """
 
     @classmethod
-    def __descendants__(cls) -> tuple[type, ...]:
-        reg: WeakSet[type] = getattr(cls, RegistryMeta._REG_ATTR)  # noqa: SLF001
+    def __descendants__(cls) -> tuple[type[Self], ...]:
+        reg: WeakSet[type[Self]] = getattr(cls, RegistryMeta._REG_ATTR)  # noqa: SLF001
         return tuple(reg)