lalamo 0.2.7__py3-none-any.whl → 0.3.1__py3-none-any.whl

lalamo/modules/linear.py

@@ -1,19 +1,25 @@
 import math
 from abc import abstractmethod
-from collections.abc import Sequence
-from dataclasses import dataclass
-from typing import NamedTuple
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass, replace
+from typing import NamedTuple, Self
 
 import equinox as eqx
 import jax
+import jax.numpy as jnp
 from einops import rearrange
-from jax import numpy as jnp
 from jaxtyping import Array, DTypeLike, Float, Int, PRNGKeyArray
 
-from lalamo.common import ParameterDict
+from lalamo.common import ParameterTree, dummy_array
 from lalamo.quantization import QuantizationMode, dynamically_quantize_activations, quantize_weights
 
-from .common import LalamoModule, WeightLayout, register_config_union
+from .common import (
+    LalamoModule,
+    WeightLayout,
+    from_layout,
+    into_layout,
+    register_config_union,
+)
 
 __all__ = [
     "FullPrecisionLinear",
@@ -48,30 +54,11 @@ class LinearBase[ConfigT: LinearConfigBase](LalamoModule[ConfigT]):
         inputs: Float[Array, " in_channels"],
     ) -> tuple[Float[Array, " out_channels"], ...]: ...
 
-    @classmethod
-    def _default_weight_layout(cls) -> WeightLayout:
-        return WeightLayout.INPUT_OUTPUT
-
-    @classmethod
-    def _into_layout(
-        cls,
-        weights: Float[Array, "in_channels out_channels"],
-        layout: WeightLayout,
-    ) -> Float[Array, "in_channels out_channels"] | Float[Array, "out_channels in_channels"]:
-        if layout == WeightLayout.AUTO:
-            layout = cls._default_weight_layout()
-        match layout:
-            case WeightLayout.OUTPUT_INPUT:
-                return weights
-            case WeightLayout.INPUT_OUTPUT:
-                return rearrange(
-                    weights,
-                    "total_out_channels in_channels -> in_channels total_out_channels",
-                )
-        raise ValueError(f"Unsupported weight layout: {layout}")
+    def __post_init__(self) -> None:
+        assert isinstance(self.output_dims, tuple)
 
-    @classmethod
-    def _get_split_points(cls, output_dims: Sequence[int]) -> tuple[int, ...]:
+    @staticmethod
+    def _get_split_points(output_dims: Sequence[int]) -> tuple[int, ...]:
         result = []
         last_split_point = 0
         for dim in output_dims[:-1]:
@@ -92,6 +79,14 @@ class LinearConfigBase:
         key: PRNGKeyArray,
     ) -> LinearBase: ...
 
+    @abstractmethod
+    def empty(
+        self,
+        input_dim: int,
+        output_dims: tuple[int, ...],
+        has_biases: bool,
+    ) -> LinearBase: ...
+
 
 @dataclass(frozen=True)
 class FullPrecisionLinearConfig(LinearConfigBase):
@@ -104,7 +99,7 @@ class FullPrecisionLinearConfig(LinearConfigBase):
         has_biases: bool,
         *,
         key: PRNGKeyArray,
-    ) -> LinearBase:
+    ) -> "FullPrecisionLinear":
         scale = 1 / math.sqrt(input_dim)
         weights = jax.random.uniform(
             key,
@@ -125,6 +120,28 @@ class FullPrecisionLinearConfig(LinearConfigBase):
             biases=biases,
         )
 
+    def empty(
+        self,
+        input_dim: int,
+        output_dims: tuple[int, ...],
+        has_biases: bool,
+    ) -> "FullPrecisionLinear":
+        weights = dummy_array(
+            (sum(output_dims), input_dim),
+            dtype=self.precision,
+        )
+        if has_biases:
+            biases = dummy_array((sum(output_dims),), dtype=self.precision)
+        else:
+            biases = None
+
+        return FullPrecisionLinear(
+            config=self,
+            output_dims=output_dims,
+            weights=weights,
+            biases=biases,
+        )
+
 
 class FullPrecisionLinear(LinearBase[FullPrecisionLinearConfig]):
     weights: Float[Array, "total_out_channels in_channels"]
@@ -148,7 +165,7 @@ class FullPrecisionLinear(LinearBase[FullPrecisionLinearConfig]):
             raise ValueError(
                 f"Weight dtype ({self.weights.dtype}) is not equal to specified precision ({self.config.precision}).",
             )
-        w_output_dim, w_input_dim = self.weights.shape
+        w_output_dim, _ = self.weights.shape
         if w_output_dim != sum(self.output_dims):
             raise ValueError(
                 f"Number of output channels in weights ({w_output_dim}) is not"
@@ -167,18 +184,31 @@ class FullPrecisionLinear(LinearBase[FullPrecisionLinearConfig]):
                 f"Bias dtype ({self.biases.dtype}) is not equal to specified precision ({self.config.precision}).",
             )
 
+    @eqx.filter_jit
     def __call__(self, inputs: Float[Array, " in_channels"]) -> tuple[Float[Array, " out_channels"], ...]:
         result = self.weights @ inputs
         if self.biases is not None:
             result = result + self.biases
         return tuple(jnp.split(result, self._get_split_points(self.output_dims)))
 
-    def export_weights(self, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterDict:
-        result = ParameterDict(weights=self._into_layout(self.weights, weight_layout))
+    def export_weights(self, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterTree:
+        result = dict(weights=into_layout(self.weights, weight_layout))
         if self.biases is not None:
             result["biases"] = self.biases
         return result
 
+    def import_weights(
+        self,
+        weights: ParameterTree[Array],
+        weight_layout: WeightLayout = WeightLayout.AUTO,
+    ) -> Self:
+        assert isinstance(weights, Mapping)
+        return replace(
+            self,
+            weights=from_layout(weights["weights"], weight_layout),
+            biases=weights["biases"] if self.has_biases else None,
+        )
+
 
 @dataclass(frozen=True)
 class GroupQuantizedLinearConfig(LinearConfigBase):
@@ -224,6 +254,34 @@ class GroupQuantizedLinearConfig(LinearConfigBase):
             biases=biases,
         )
 
+    def empty(
+        self,
+        input_dim: int,
+        output_dims: tuple[int, ...],
+        has_biases: bool,
+    ) -> LinearBase:
+        weights = dummy_array(
+            (sum(output_dims), input_dim),
+            dtype=self.activation_precision,
+        )
+        num_groups = input_dim // self.group_size
+        scales = dummy_array((sum(output_dims), num_groups), dtype=self.activation_precision)
+
+        if has_biases:
+            biases = dummy_array((sum(output_dims),), dtype=self.activation_precision)
+        else:
+            biases = None
+        zero_points = dummy_array((sum(output_dims), num_groups), dtype=self.activation_precision)
+
+        return GroupQuantizedLinear(
+            config=self,
+            output_dims=output_dims,
+            weights=weights,
+            scales=scales,
+            zero_points=zero_points,
+            biases=biases,
+        )
+
 
 class RequantizedWeights(NamedTuple):
     weights: Int[Array, "total_out_channels in_channels"]
@@ -271,7 +329,7 @@ class GroupQuantizedLinearBase[ConfigT: GroupQuantizedLinearConfig](LinearBase[C
                 f" ({self.config.activation_precision}).",
                 " Quantized layers require parameter dtypes to be equal to the activation precision.",
             )
-        w_output_dim, w_input_dim = self.weights.shape
+        w_output_dim, _ = self.weights.shape
         if w_output_dim != sum(self.output_dims):
             raise ValueError(
                 f"Number of output channels in weights ({w_output_dim}) is not"
@@ -352,100 +410,20 @@ class GroupQuantizedLinearBase[ConfigT: GroupQuantizedLinearConfig](LinearBase[C
             inputs = dynamically_quantize_activations(inputs, self.config.activation_quantization_mode)
         return self._prepare_scaled_weights() @ inputs
 
+    @eqx.filter_jit
     def __call__(self, inputs: Float[Array, " in_channels"]) -> tuple[Float[Array, " out_channels"], ...]:
         result = self._apply_weights(inputs)
         if self.biases is not None:
             result = result + self.biases
         return tuple(jnp.split(result, self._get_split_points(self.output_dims)))
 
-    def requantize_weights(self, weights, zero_points, scales):
-        """
-        Requantize weights from [20, 6144] grouping to [2560, 48] grouping.
-
-        Args:
-            weights: uint4 array of shape [M, N]
-            zero_points: uint4 array of shape [M//group_size_0, N//group_size_1]
-            scales: float16 array of shape [M//group_size_0, N//group_size_1]
-
-        Returns:
-            new_weights: uint4 array of shape [M, N]
-            new_zero_points: uint4 array of shape [M, N//128]
-            new_scales: float16 array of shape [M, N//128]
-        """
-        # Get dimensions
-        M, N = weights.shape
-        old_groups_0, old_groups_1 = zero_points.shape
-
-        # Calculate old group sizes
-        old_group_size_0 = M // old_groups_0  # 2560 // 20 = 128
-        old_group_size_1 = N // old_groups_1  # 6144 // 6144 = 1
-
-        # New group sizes
-        new_group_size_0 = 1  # 2560 // 2560 = 1
-        new_group_size_1 = self.config.group_size  # 6144 // 48 = 128
-
-        # Step 1: Dequantize with original parameters
-        # Expand zero_points and scales to match weights shape
-        zp_expanded = jnp.repeat(jnp.repeat(zero_points, old_group_size_0, axis=0), old_group_size_1, axis=1)
-        scales_expanded = jnp.repeat(jnp.repeat(scales, old_group_size_0, axis=0), old_group_size_1, axis=1)
-
-        # Dequantize (convert to float for computation)
-        weights_float = weights.astype(jnp.float32)
-        zp_float = zp_expanded.astype(jnp.float32)
-        dequantized = (weights_float - zp_float) * scales_expanded.astype(jnp.float32)
-
-        # Step 2: Requantize with new group structure [2560, 48]
-        # Reshape for new groups
-        dequantized_reshaped = dequantized.reshape(
-            M // new_group_size_0,
-            new_group_size_0,
-            N // new_group_size_1,
-            new_group_size_1,
-        )
-
-        # Compute new scales and zero points per group
-        # Move group dimensions to the end for reduction
-        dequantized_groups = dequantized_reshaped.transpose(0, 2, 1, 3)  # [2560, 48, 1, 128]
-
-        # Find min and max per group
-        group_min = dequantized_groups.min(axis=(2, 3), keepdims=True)
-        group_max = dequantized_groups.max(axis=(2, 3), keepdims=True)
-
-        # Compute scales (with small epsilon to avoid division by zero)
-        eps = 1e-6
-        new_scales = ((group_max - group_min) / 15.0 + eps).astype(scales.dtype)
-        new_scales = new_scales.squeeze(axis=(2, 3))  # [2560, 48]
-
-        # Compute zero points (quantize to uint4 range 0-15)
-        new_zero_points = jnp.round(-group_min.squeeze(axis=(2, 3)) / new_scales).astype(jnp.uint4)
-        new_zero_points = jnp.clip(new_zero_points, 0, 15)
-
-        # Quantize with new parameters
-        scales_expanded_new = jnp.repeat(new_scales, new_group_size_1, axis=1).reshape(M, N)
-        zp_expanded_new = jnp.repeat(new_zero_points, new_group_size_1, axis=1).reshape(M, N)
-
-        new_weights = jnp.round(
-            dequantized / scales_expanded_new.astype(jnp.float32) + zp_expanded_new.astype(jnp.float32),
-        )
-        new_weights = jnp.clip(new_weights, 0, 15).astype(jnp.uint4)
-
-        return new_weights, new_zero_points, new_scales
-
-    def export_weights(self, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterDict:
-        exported_weights = self._into_layout(self.int_weights, weight_layout)
-
-        exported_zero_points = self._into_layout(self.int_zero_points, weight_layout)
+    def export_weights(self, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterTree:
+        expected_weight_layout = WeightLayout.OUTPUT_INPUT
+        exported_weights = into_layout(self.int_weights, expected_weight_layout)
+        exported_zero_points = into_layout(self.int_zero_points, expected_weight_layout)
+        exported_scales = into_layout(self.scales, expected_weight_layout)
 
-        exported_scales = self._into_layout(self.scales, weight_layout)
-
-        # CRIMINAL HACK!!!
-        exported_weights, exported_zero_points, exported_scales = self.requantize_weights(
-            exported_weights,
-            exported_zero_points,
-            exported_scales,
-        )
-
-        result = ParameterDict(
+        result = dict(
             weights=exported_weights,
             zero_points=exported_zero_points,
             scales=exported_scales,
@@ -454,6 +432,21 @@ class GroupQuantizedLinearBase[ConfigT: GroupQuantizedLinearConfig](LinearBase[C
             result["biases"] = self.biases
         return result
 
+    def import_weights(
+        self,
+        weights: ParameterTree[Array],
+        weight_layout: WeightLayout = WeightLayout.AUTO,
+    ) -> Self:
+        assert isinstance(weights, Mapping)
+        assert isinstance(weights["weights"], Array)
+        return replace(
+            self,
+            weights=from_layout(weights["weights"].astype(self.weights.dtype), weight_layout),
+            scales=from_layout(weights["scales"], weight_layout),
+            zero_points=from_layout(weights["zero_points"], weight_layout).astype(self.zero_points.dtype),
+            biases=weights["biases"] if self.has_biases else None,
+        )
+
 
 class GroupQuantizedLinear(GroupQuantizedLinearBase[GroupQuantizedLinearConfig]):
     pass
@@ -512,6 +505,38 @@ class QLoRALinearConfig(GroupQuantizedLinearConfig):
             lora_up_weights=lora_up_weights,
         )
 
+    def empty(
+        self,
+        input_dim: int,
+        output_dims: tuple[int, ...],
+        has_biases: bool,
+    ) -> LinearBase:
+        group_quantized_linear = super().empty(input_dim, output_dims, has_biases)
+        assert isinstance(group_quantized_linear, GroupQuantizedLinear)
+        hidden_lora_rank = len(output_dims) * self.lora_rank
+        lora_down_weights = dummy_array(
+            (hidden_lora_rank, input_dim),
+            dtype=self.activation_precision,
+        )
+        lora_up_weights = tuple(
+            dummy_array(
+                (output_dim, self.lora_rank),
+                dtype=self.activation_precision,
+            )
+            for output_dim in output_dims
+        )
+
+        return QLoRALinear(
+            config=self,
+            output_dims=output_dims,
+            weights=group_quantized_linear.weights,
+            scales=group_quantized_linear.scales,
+            biases=group_quantized_linear.biases,
+            zero_points=group_quantized_linear.zero_points,
+            lora_down_weights=lora_down_weights,
+            lora_up_weights=lora_up_weights,
+        )
+
 
 class QLoRALinear(GroupQuantizedLinearBase[QLoRALinearConfig]):
     lora_down_weights: Float[Array, "total_lora_channels in_channels"]
@@ -564,6 +589,7 @@ class QLoRALinear(GroupQuantizedLinearBase[QLoRALinearConfig]):
                     f" equal to lora_rank ({self.config.lora_rank}).",
                 )
 
+    @eqx.filter_jit
     def __call__(self, inputs: Float[Array, " in_channels"]) -> tuple[Float[Array, " out_channels"], ...]:
         joint_q_out = self._apply_weights(inputs)
         q_outs = jnp.split(joint_q_out, self._get_split_points(self.output_dims))
@@ -584,16 +610,30 @@ class QLoRALinear(GroupQuantizedLinearBase[QLoRALinearConfig]):
 
         return tuple(results)
 
-    def export_weights(self, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterDict:
-        quantized_linear_weights = super().export_weights()
-        exported_lora_down_weights = self._into_layout(self.lora_down_weights, weight_layout)
-        exported_lora_up_weights = tuple(
-            self._into_layout(lora_up_weight, weight_layout) for lora_up_weight in self.lora_up_weights
-        )
-        return ParameterDict(
+    def export_weights(self, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterTree:
+        quantized_linear_weights: dict[str, ParameterTree] = super().export_weights()  # type: ignore
+        exported_lora_down_weights = into_layout(self.lora_down_weights, weight_layout)
+        exported_lora_up_weights = [
+            into_layout(lora_up_weight, weight_layout) for lora_up_weight in self.lora_up_weights
+        ]
+        return dict(
+            down_weights=into_layout(exported_lora_down_weights, weight_layout),
+            up_weights=[into_layout(w, weight_layout) for w in exported_lora_up_weights],
             **quantized_linear_weights,
-            down_weights=exported_lora_down_weights,
-            up_weights=exported_lora_up_weights,
+        )
+
+    def import_weights(
+        self,
+        weights: ParameterTree[Array],
+        weight_layout: WeightLayout = WeightLayout.AUTO,
+    ) -> Self:
+        base = super().import_weights(weights, weight_layout)
+        assert isinstance(weights, Mapping)
+        assert isinstance(weights["up_weights"], Sequence)
+        return replace(
+            base,
+            lora_down_weights=from_layout(weights["down_weights"], weight_layout),
+            lora_up_weights=tuple(from_layout(up_weights, weight_layout) for up_weights in weights["up_weights"]),
         )
 
 

lalamo/modules/mlp.py

@@ -1,9 +1,12 @@
-from dataclasses import dataclass
+from collections.abc import Mapping
+from dataclasses import dataclass, replace
+from typing import Self
 
+import equinox as eqx
 import jax
 from jaxtyping import Array, DTypeLike, Float, PRNGKeyArray
 
-from lalamo.common import ParameterDict
+from lalamo.common import ParameterTree
 
 from .activations import Activation
 from .common import LalamoModule, WeightLayout
@@ -35,8 +38,23 @@ class MLPConfig:
             ),
         )
 
+    def empty(self, model_dim: int, hidden_dim: int) -> "MLP":
+        return MLP(
+            self,
+            up_projection=self.linear_config.empty(
+                model_dim,
+                (hidden_dim, hidden_dim),
+                has_biases=False,
+            ),
+            down_projection=self.linear_config.empty(
+                hidden_dim,
+                (model_dim,),
+                has_biases=False,
+            ),
+        )
 
-class MLP(LalamoModule):
+
+class MLP(LalamoModule[MLPConfig]):
     up_projection: LinearBase
     down_projection: LinearBase
 
@@ -66,14 +84,29 @@ class MLP(LalamoModule):
                 f" the up projection output dimension {self.up_projection.input_dim}",
             )
 
+    @eqx.filter_jit
     def __call__(self, inputs: Float[Array, " channels"]) -> Float[Array, " channels"]:
         up_proj, gate = self.up_projection(inputs)
         gate = self.config.activation(gate)
         (result,) = self.down_projection(up_proj * gate)
         return result
 
-    def export_weights(self, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterDict:
-        return ParameterDict(
-            up_projection=self.up_projection.export_weights(weight_layout),
-            down_projection=self.down_projection.export_weights(weight_layout),
+    def export_weights(self, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterTree:
+        return {
+            "up_projection": self.up_projection.export_weights(weight_layout),
+            "down_projection": self.down_projection.export_weights(weight_layout),
+        }
+
+    def import_weights(
+        self,
+        weights: ParameterTree[Array],
+        weight_layout: WeightLayout = WeightLayout.AUTO,
+    ) -> 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"], weight_layout),
+            down_projection=self.down_projection.import_weights(weights["down_projection"], weight_layout),
         )

lalamo/modules/normalization.py

@@ -1,11 +1,14 @@
-from dataclasses import dataclass
+from collections.abc import Mapping
+from dataclasses import dataclass, replace
 from enum import Enum
+from typing import Self
 
+import equinox as eqx
 import jax
 from jax import numpy as jnp
 from jaxtyping import Array, DTypeLike, Float
 
-from lalamo.common import ParameterDict
+from lalamo.common import ParameterTree, dummy_array
 
 from .common import LalamoModule, WeightLayout
 
@@ -29,10 +32,16 @@ class RMSNormConfig:
     scale_offset: float | None
     upcast_mode: UpcastMode
 
-    def init(self, channels: int) -> "RMSNorm":
-        scales = jnp.ones(channels, dtype=self.scale_precision)
+    def init(self, input_dim: int) -> "RMSNorm":
+        scales = jnp.ones(input_dim, dtype=self.scale_precision)
         return RMSNorm(self, scales=scales)
 
+    def empty(self, input_dim: int) -> "RMSNorm":
+        return RMSNorm(
+            config=self,
+            scales=dummy_array(input_dim, dtype=self.scale_precision),
+        )
+
 
 class RMSNorm(LalamoModule[RMSNormConfig]):
     scales: Float[Array, " channels"]
@@ -53,6 +62,7 @@ class RMSNorm(LalamoModule[RMSNormConfig]):
                 f" specified precision {self.config.scale_precision}",
             )
 
+    @eqx.filter_jit
     def __call__(self, inputs: Float[Array, " channels"]) -> Float[Array, " channels"]:
         upcasted_inputs = inputs.astype(self.config.accumulation_precision)
 
@@ -73,5 +83,13 @@ class RMSNorm(LalamoModule[RMSNormConfig]):
         result = normalized_x * adjusted_scales
         return result.astype(inputs.dtype)
 
-    def export_weights(self, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterDict:  # noqa: ARG002
-        return ParameterDict(scales=self.scales)
+    def export_weights(self, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterTree:  # noqa: ARG002
+        return {"scales": self.scales}
+
+    def import_weights(
+        self,
+        weights: ParameterTree[Array],
+        weight_layout: WeightLayout = WeightLayout.AUTO,  # noqa: ARG002
+    ) -> Self:
+        assert isinstance(weights, Mapping)
+        return replace(self, scales=weights["scales"])

lalamo/modules/rope.py

@@ -16,13 +16,14 @@
 # limitations under the License.
 
 import math
-from dataclasses import dataclass
+from collections.abc import Mapping
+from dataclasses import dataclass, replace
 
 import equinox as eqx
 from jax import numpy as jnp
 from jaxtyping import Array, DTypeLike, Float, Int
 
-from lalamo.common import ParameterDict
+from lalamo.common import ParameterTree
 
 from .common import LalamoModule, WeightLayout, register_config_union
 
@@ -53,8 +54,8 @@ class PositionalEmbeddings(eqx.Module):
     def apply(self, heads: Float[Array, "tokens head_channels"]) -> Float[Array, "tokens head_channels"]:
         return heads * self.cosines + self.rotate_half(heads) * self.sines
 
-    def export(self, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterDict:  # noqa: ARG002
-        return ParameterDict(
+    def export(self, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterTree:  # noqa: ARG002
+        return dict(
             cosines=self.cosines,
             sines=self.sines,
         )
@@ -103,6 +104,11 @@ class RoPE(LalamoModule[RoPEConfigBase]):
         return self.config.precision
 
     def __post_init__(self) -> None:
+        num_tokens, _ = self.sines.shape
+        if num_tokens != self.config.max_sequence_length:
+            raise ValueError(
+                f"{num_tokens} does not match the specified max sequence length {self.config.max_sequence_length}",
+            )
         if self.cosines.dtype != self.config.precision:
             raise ValueError(
                 f"Cosines dtype {self.cosines.dtype} does not match the specified precision {self.config.precision}",
@@ -127,14 +133,26 @@ class RoPE(LalamoModule[RoPEConfigBase]):
         result, _ = self.sines.shape
         return result
 
+    @eqx.filter_jit
     def __call__(self, timesteps: Int[Array, " tokens"]) -> PositionalEmbeddings:
         return PositionalEmbeddings(
             cosines=self.cosines[timesteps],
             sines=self.sines[timesteps],
         )
 
-    def export_weights(self, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterDict:  # noqa: ARG002
-        return ParameterDict(cosines=self.cosines, sines=self.sines)
+    def export_weights(self, weight_layout: WeightLayout = WeightLayout.AUTO) -> ParameterTree[Array]:  # noqa: ARG002
+        return {
+            "cosines": self.cosines,
+            "sines": self.sines,
+        }
+
+    def import_weights(
+        self,
+        weights: ParameterTree[Array],
+        weight_layout: WeightLayout = WeightLayout.AUTO,  # noqa: ARG002
+    ) -> "RoPE":
+        assert isinstance(weights, Mapping)
+        return replace(self, cosines=weights["cosines"], sines=weights["sines"])
 
 
 class UnscaledRoPEConfig(RoPEConfigBase):