keras-hub-nightly 0.22.0.dev202507150421__py3-none-any.whl → 0.22.0.dev202507170424__py3-none-any.whl

keras_hub/src/models/dinov2/dinov2_presets.py

@@ -0,0 +1,4 @@
+"""DINOV2 model preset configurations."""
+
+# Metadata for loading pretrained model weights.
+backbone_presets = {}

keras_hub/src/models/flux/flux_text_to_image_preprocessor.py

@@ -43,9 +43,13 @@ class FluxTextToImagePreprocessor(Preprocessor):
 
     def generate_preprocess(self, x):
         token_ids = {}
-        token_ids["clip_l"] = self.clip_l_preprocessor(x)["token_ids"]
+        token_ids["clip_l"] = self.clip_l_preprocessor(
+            {"prompts": x, "images": None}
+        )["token_ids"]
         if self.t5_preprocessor is not None:
-            token_ids["t5"] = self.t5_preprocessor(x)["token_ids"]
+            token_ids["t5"] = self.t5_preprocessor(
+                {"prompts": x, "images": None}
+            )["token_ids"]
         return token_ids
 
     def get_config(self):

keras_hub/src/models/hgnetv2/__init__.py

@@ -0,0 +1,5 @@
+from keras_hub.src.models.hgnetv2.hgnetv2_backbone import HGNetV2Backbone
+from keras_hub.src.models.hgnetv2.hgnetv2_presets import backbone_presets
+from keras_hub.src.utils.preset_utils import register_presets
+
+register_presets(backbone_presets, HGNetV2Backbone)

keras_hub/src/models/hgnetv2/hgnetv2_presets.py

@@ -9,7 +9,7 @@ backbone_presets = {
             "params": 13599072,
             "path": "hgnetv2",
         },
-        "kaggle_handle": "",
+        "kaggle_handle": "kaggle://keras/hgnetv2/keras/hgnetv2_b4_ssld_stage2_ft_in1k/1",
     },
     "hgnetv2_b5_ssld_stage1_in22k_in1k": {
         "metadata": {
@@ -20,7 +20,7 @@ backbone_presets = {
             "params": 33419680,
             "path": "hgnetv2",
         },
-        "kaggle_handle": "",
+        "kaggle_handle": "kaggle://keras/hgnetv2/keras/hgnetv2_b5_ssld_stage1_in22k_in1k/1",
     },
     "hgnetv2_b5_ssld_stage2_ft_in1k": {
         "metadata": {
@@ -31,7 +31,7 @@ backbone_presets = {
             "params": 33419680,
             "path": "hgnetv2",
         },
-        "kaggle_handle": "",
+        "kaggle_handle": "kaggle://keras/hgnetv2/keras/hgnetv2_b5_ssld_stage2_ft_in1k/1",
     },
     "hgnetv2_b6_ssld_stage1_in22k_in1k": {
         "metadata": {
@@ -42,7 +42,7 @@ backbone_presets = {
             "params": 69179888,
             "path": "hgnetv2",
         },
-        "kaggle_handle": "",
+        "kaggle_handle": "kaggle://keras/hgnetv2/keras/hgnetv2_b6_ssld_stage1_in22k_in1k/1",
     },
     "hgnetv2_b6_ssld_stage2_ft_in1k": {
         "metadata": {
@@ -53,6 +53,6 @@ backbone_presets = {
             "params": 69179888,
             "path": "hgnetv2",
         },
-        "kaggle_handle": "",
+        "kaggle_handle": "kaggle://keras/hgnetv2/keras/hgnetv2_b6_ssld_stage2_ft_in1k/1",
     },
 }

keras_hub/src/models/stable_diffusion_3/flow_match_euler_discrete_scheduler.py

@@ -38,7 +38,6 @@ class FlowMatchEulerDiscreteScheduler(layers.Layer):
         timesteps = ops.flip(timesteps, axis=0)
         sigmas = self._timestep_to_sigma(timesteps)
 
-        self.timesteps = ops.multiply(sigmas, num_train_timesteps)
         self.sigma_min = sigmas[-1]
         self.sigma_max = sigmas[0]
 
@@ -54,14 +53,24 @@ class FlowMatchEulerDiscreteScheduler(layers.Layer):
             )
         return sigma
 
+    def set_sigmas(self, num_steps):
+        timesteps = ops.linspace(
+            self._sigma_to_timestep(self.sigma_max),
+            self._sigma_to_timestep(self.sigma_min),
+            num_steps,
+        )
+        sigmas = self._timestep_to_sigma(timesteps)
+        sigmas = ops.concatenate([sigmas, ops.zeros((1,), dtype=sigmas.dtype)])
+        self.sigmas = sigmas
+
     def call(self, inputs, num_steps):
-        start = self._sigma_to_timestep(self.sigma_max)
-        end = self._sigma_to_timestep(self.sigma_min)
-        step_size = ops.divide(
-            ops.subtract(end, start), ops.subtract(num_steps, 1)
+        if not hasattr(self, "sigmas"):
+            self.set_sigmas(num_steps)
+
+        step = ops.expand_dims(
+            ops.convert_to_tensor(inputs, dtype="int32"), axis=0
         )
-        timestep = ops.add(start, ops.multiply(inputs, step_size))
-        sigma = ops.maximum(self._timestep_to_sigma(timestep), 0.0)
+        sigma = ops.take(self.sigmas, step)
         timestep = self._sigma_to_timestep(sigma)
         return sigma, timestep
 

keras_hub/src/models/stable_diffusion_3/mmdit.py

@@ -10,6 +10,63 @@ from keras_hub.src.utils.keras_utils import fused_attention_op_available
 from keras_hub.src.utils.keras_utils import gelu_approximate
 from keras_hub.src.utils.keras_utils import standardize_data_format
 
+# TODO: Deprecate this in favor of
+# `keras.layers.RMSNormalization` once we require Keras 3.9 or later.
+if hasattr(layers, "RMSNormalization"):
+    RMSNormalization = layers.RMSNormalization
+else:
+
+    class RMSNormalization(layers.Layer):
+        """A normalization layer for MMDiT that implements RMS normalization."""
+
+        def __init__(self, axis=-1, epsilon=1e-6, **kwargs):
+            super().__init__(**kwargs)
+            self.axis = axis
+            self.epsilon = epsilon
+
+        def build(self, input_shape):
+            if isinstance(self.axis, list):
+                shape = tuple([input_shape[dim] for dim in self.axis])
+            else:
+                shape = (input_shape[self.axis],)
+                self.axis = [self.axis]
+
+            self.scale = self.add_weight(
+                name="scale", shape=shape, initializer="ones"
+            )
+
+            self.built = True
+
+        def call(self, x):
+            x = ops.cast(
+                x, keras.backend.result_type(self.compute_dtype, "float32")
+            )
+            rrms = ops.rsqrt(
+                ops.mean(ops.square(x), axis=self.axis, keepdims=True)
+                + self.epsilon
+            )
+            return (x * rrms) * ops.cast(self.scale, x.dtype)
+
+        def compute_output_shape(self, input_shape):
+            if isinstance(self.axis, int):
+                axes = [self.axis]
+            else:
+                axes = self.axis
+
+            for axis in axes:
+                if axis >= len(input_shape) or axis < -len(input_shape):
+                    raise ValueError(
+                        f"Axis {axis} is out of bounds for "
+                        f"input shape {input_shape}. "
+                        f"Received: axis={self.axis}"
+                    )
+            return input_shape
+
+        def get_config(self):
+            config = super().get_config()
+            config.update({"axis": self.axis, "epsilon": self.epsilon})
+            return config
+
 
 class AdaptiveLayerNormalization(layers.Layer):
     """Adaptive layer normalization.
@@ -402,11 +459,11 @@ def get_qk_norm(qk_norm=None, q_norm_name="q_norm", k_norm_name="k_norm"):
     if qk_norm is None:
         pass
     elif qk_norm == "rms_norm":
-        q_norm = layers.LayerNormalization(
-            epsilon=1e-6, rms_scaling=True, dtype="float32", name=q_norm_name
+        q_norm = RMSNormalization(
+            axis=-1, epsilon=1e-6, dtype="float32", name=q_norm_name
         )
-        k_norm = layers.LayerNormalization(
-            epsilon=1e-6, rms_scaling=True, dtype="float32", name=k_norm_name
+        k_norm = RMSNormalization(
+            axis=-1, epsilon=1e-6, dtype="float32", name=k_norm_name
         )
     else:
         raise NotImplementedError(

keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_backbone.py

@@ -96,26 +96,10 @@ class LatentRescaling(layers.Rescaling):
         return (self.backend.cast(inputs, dtype) / scale) + offset
 
 
-class ClassifierFreeGuidanceConcatenate(layers.Layer):
-    def call(
-        self,
-        latents,
-        positive_contexts,
-        negative_contexts,
-        positive_pooled_projections,
-        negative_pooled_projections,
-        timestep,
-    ):
+class TimestepBroadcastTo(layers.Layer):
+    def call(self, latents, timestep):
         timestep = ops.broadcast_to(timestep, ops.shape(latents)[:1])
-        latents = ops.concatenate([latents, latents], axis=0)
-        contexts = ops.concatenate(
-            [positive_contexts, negative_contexts], axis=0
-        )
-        pooled_projections = ops.concatenate(
-            [positive_pooled_projections, negative_pooled_projections], axis=0
-        )
-        timesteps = ops.concatenate([timestep, timestep], axis=0)
-        return latents, contexts, pooled_projections, timesteps
+        return timestep
 
 
 class ClassifierFreeGuidance(layers.Layer):
@@ -330,8 +314,8 @@ class StableDiffusion3Backbone(Backbone):
             name="diffuser",
         )
         self.vae = vae
-        self.cfg_concat = ClassifierFreeGuidanceConcatenate(
-            dtype=dtype, name="classifier_free_guidance_concat"
+        self.timestep_broadcast_to = TimestepBroadcastTo(
+            dtype=dtype, name="timestep_broadcast_to"
         )
         self.cfg = ClassifierFreeGuidance(
             dtype=dtype, name="classifier_free_guidance"
@@ -538,6 +522,9 @@ class StableDiffusion3Backbone(Backbone):
         latents = self.vae.encode(images)
         return self.image_rescaling(latents)
 
+    def configure_scheduler(self, num_steps):
+        self.scheduler.set_sigmas(num_steps)
+
     def add_noise_step(self, latents, noises, step, num_steps):
         return self.scheduler.add_noise(latents, noises, step, num_steps)
 
@@ -562,11 +549,15 @@ class StableDiffusion3Backbone(Backbone):
 
         # Concatenation for classifier-free guidance.
         if guidance_scale is not None:
-            concated_latents, contexts, pooled_projs, timesteps = (
-                self.cfg_concat(latents, *embeddings, timestep)
+            timestep = self.timestep_broadcast_to(latents, timestep)
+            timesteps = ops.concatenate([timestep, timestep], axis=0)
+            concated_latents = ops.concatenate([latents, latents], axis=0)
+            contexts = ops.concatenate([embeddings[0], embeddings[1]], axis=0)
+            pooled_projs = ops.concatenate(
+                [embeddings[2], embeddings[3]], axis=0
             )
         else:
-            timesteps = ops.broadcast_to(timestep, ops.shape(latents)[:1])
+            timesteps = self.timestep_broadcast_to(latents, timestep)
             concated_latents = latents
             contexts = embeddings[0]
             pooled_projs = embeddings[2]
@@ -623,20 +614,20 @@ class StableDiffusion3Backbone(Backbone):
     def from_config(cls, config, custom_objects=None):
         config = config.copy()
 
-        # Propagate `dtype` to text encoders if needed.
+        # Propagate `dtype` to the VAE if needed.
         if "dtype" in config and config["dtype"] is not None:
             dtype_config = config["dtype"]
             if "dtype" not in config["vae"]["config"]:
                 config["vae"]["config"]["dtype"] = dtype_config
-            if "dtype" not in config["clip_l"]["config"]:
-                config["clip_l"]["config"]["dtype"] = dtype_config
-            if "dtype" not in config["clip_g"]["config"]:
-                config["clip_g"]["config"]["dtype"] = dtype_config
+
+        # Text encoders default to float16 dtype if not specified.
+        for text_encoder in ("clip_l", "clip_g", "t5"):
             if (
-                config["t5"] is not None
-                and "dtype" not in config["t5"]["config"]
+                text_encoder in config
+                and config[text_encoder] is not None
+                and "dtype" not in config[text_encoder]["config"]
             ):
-                config["t5"]["config"]["dtype"] = dtype_config
+                config[text_encoder]["config"]["dtype"] = "float16"
 
         # We expect `vae`, `clip_l`, `clip_g` and/or `t5` to be instantiated.
         config["vae"] = layers.deserialize(

keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_image_to_image.py

@@ -169,6 +169,7 @@ class StableDiffusion3ImageToImage(ImageToImage):
         guidance_scale=7.0,
         seed=None,
     ):
+        self.backbone.configure_scheduler(num_steps)
         return super().generate(
             inputs,
             num_steps=num_steps,

keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_inpaint.py

@@ -184,6 +184,7 @@ class StableDiffusion3Inpaint(Inpaint):
         guidance_scale=7.0,
         seed=None,
     ):
+        self.backbone.configure_scheduler(num_steps)
         return super().generate(
             inputs,
             num_steps=num_steps,

keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_text_to_image.py

@@ -141,6 +141,7 @@ class StableDiffusion3TextToImage(TextToImage):
         guidance_scale=7.0,
         seed=None,
     ):
+        self.backbone.configure_scheduler(num_steps)
         return super().generate(
             inputs,
             num_steps=num_steps,

keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_text_to_image_preprocessor.py

@@ -50,8 +50,12 @@ class StableDiffusion3TextToImagePreprocessor(TextToImagePreprocessor):
 
     def generate_preprocess(self, x):
         token_ids = {}
-        token_ids["clip_l"] = self.clip_l_preprocessor(x)["token_ids"]
-        token_ids["clip_g"] = self.clip_g_preprocessor(x)["token_ids"]
+        token_ids["clip_l"] = self.clip_l_preprocessor(
+            {"prompts": x, "images": None}
+        )["token_ids"]
+        token_ids["clip_g"] = self.clip_g_preprocessor(
+            {"prompts": x, "images": None}
+        )["token_ids"]
         if self.t5_preprocessor is not None:
             token_ids["t5"] = self.t5_preprocessor(x)["token_ids"]
         return token_ids

keras_hub/src/utils/preset_utils.py

@@ -649,7 +649,10 @@ class KerasPresetLoader(PresetLoader):
         return check_config_class(self.config)
 
     def load_backbone(self, cls, load_weights, **kwargs):
-        backbone = self._load_serialized_object(self.config, **kwargs)
+        config = self.config.copy()
+        backbone_kwargs, kwargs = self.get_backbone_kwargs(**kwargs)
+        config["config"] = {**config["config"], **backbone_kwargs}
+        backbone = self._load_serialized_object(config, **kwargs)
         if load_weights:
             jax_memory_cleanup(backbone)
             self._load_backbone_weights(backbone)

keras_hub/src/utils/transformers/convert_dinov2.py

@@ -0,0 +1,180 @@
+import numpy as np
+
+from keras_hub.src.models.dinov2.dinov2_backbone import DINOV2Backbone
+
+backbone_cls = DINOV2Backbone
+
+
+def convert_backbone_config(transformers_config):
+    model_type = transformers_config["model_type"]
+    antialias_in_interpolation = False if model_type == "dinov2" else True
+    image_size = transformers_config["image_size"]
+    intermediate_dim = int(
+        transformers_config["hidden_size"] * transformers_config["mlp_ratio"]
+    )
+    return {
+        "patch_size": transformers_config["patch_size"],
+        "num_layers": transformers_config["num_hidden_layers"],
+        "hidden_dim": transformers_config["hidden_size"],
+        "num_heads": transformers_config["num_attention_heads"],
+        "intermediate_dim": intermediate_dim,
+        "layer_scale_init_value": transformers_config["layerscale_value"],
+        "num_register_tokens": transformers_config.get(
+            "num_register_tokens", 0
+        ),
+        "use_mask_token": transformers_config.get("use_mask_token", True),
+        "use_swiglu_ffn": transformers_config["use_swiglu_ffn"],
+        "dropout_rate": transformers_config["hidden_dropout_prob"],
+        "drop_path_rate": transformers_config["drop_path_rate"],
+        "image_shape": (image_size, image_size, 3),
+        "position_embedding_shape": (image_size, image_size),
+        "antialias_in_interpolation": antialias_in_interpolation,
+    }
+
+
+def convert_weights(backbone, loader, transformers_config):
+    if not isinstance(backbone, DINOV2Backbone):
+        raise ValueError(
+            "The provided backbone must be an instance of DINOV2Backbone. "
+            f"Received: {type(backbone)}"
+        )
+
+    def port_ln(keras_variable, weight_key):
+        loader.port_weight(keras_variable.gamma, f"{weight_key}.weight")
+        loader.port_weight(keras_variable.beta, f"{weight_key}.bias")
+
+    def port_dense(keras_variable, weight_key):
+        loader.port_weight(
+            keras_variable.kernel,
+            f"{weight_key}.weight",
+            hook_fn=lambda x, _: x.T,
+        )
+        if keras_variable.bias is not None:
+            loader.port_weight(keras_variable.bias, f"{weight_key}.bias")
+
+    def port_mha(keras_variable, weight_key, num_heads, hidden_dim):
+        # query
+        loader.port_weight(
+            keras_variable.query_dense.kernel,
+            f"{weight_key}.attention.query.weight",
+            hook_fn=lambda x, _: np.reshape(
+                x.T, (hidden_dim, num_heads, hidden_dim // num_heads)
+            ),
+        )
+        loader.port_weight(
+            keras_variable.query_dense.bias,
+            f"{weight_key}.attention.query.bias",
+            hook_fn=lambda x, _: np.reshape(
+                x, (num_heads, hidden_dim // num_heads)
+            ),
+        )
+        # key
+        loader.port_weight(
+            keras_variable.key_dense.kernel,
+            f"{weight_key}.attention.key.weight",
+            hook_fn=lambda x, _: np.reshape(
+                x.T, (hidden_dim, num_heads, hidden_dim // num_heads)
+            ),
+        )
+        loader.port_weight(
+            keras_variable.key_dense.bias,
+            f"{weight_key}.attention.key.bias",
+            hook_fn=lambda x, _: np.reshape(
+                x, (num_heads, hidden_dim // num_heads)
+            ),
+        )
+        # value
+        loader.port_weight(
+            keras_variable.value_dense.kernel,
+            f"{weight_key}.attention.value.weight",
+            hook_fn=lambda x, _: np.reshape(
+                x.T, (hidden_dim, num_heads, hidden_dim // num_heads)
+            ),
+        )
+        loader.port_weight(
+            keras_variable.value_dense.bias,
+            f"{weight_key}.attention.value.bias",
+            hook_fn=lambda x, _: np.reshape(
+                x, (num_heads, hidden_dim // num_heads)
+            ),
+        )
+        # output
+        loader.port_weight(
+            keras_variable.output_dense.kernel,
+            f"{weight_key}.output.dense.weight",
+            hook_fn=lambda x, _: np.reshape(
+                x.T, (num_heads, hidden_dim // num_heads, hidden_dim)
+            ),
+        )
+        loader.port_weight(
+            keras_variable.output_dense.bias, f"{weight_key}.output.dense.bias"
+        )
+
+    # Embedding.
+    loader.port_weight(
+        keras_variable=backbone.embeddings.cls_token,
+        hf_weight_key="embeddings.cls_token",
+    )
+    if backbone.use_mask_token:
+        loader.port_weight(
+            keras_variable=backbone.embeddings.mask_token,
+            hf_weight_key="embeddings.mask_token",
+        )
+    if backbone.num_register_tokens > 0:
+        loader.port_weight(
+            keras_variable=backbone.embeddings.register_tokens,
+            hf_weight_key="embeddings.register_tokens",
+        )
+    loader.port_weight(
+        keras_variable=backbone.embeddings.position_embeddings,
+        hf_weight_key="embeddings.position_embeddings",
+    )
+    # Interpolate position embeddings to match the image shape.
+    backbone.embeddings.interpolated_position_embeddings.assign(
+        backbone.embeddings._interpolate_position_embeddings(
+            backbone.embeddings.position_embeddings,
+            patch_size=backbone.patch_size,
+            source_shape=backbone.embeddings.position_embedding_shape,
+            target_shape=backbone.image_shape,
+            antialias=backbone.embeddings.antialias_in_interpolation,
+        )
+    )
+    loader.port_weight(
+        keras_variable=backbone.embeddings.patch_embeddings.projection.kernel,
+        hf_weight_key="embeddings.patch_embeddings.projection.weight",
+        hook_fn=lambda x, _: np.transpose(x, (2, 3, 1, 0)),
+    )
+    loader.port_weight(
+        keras_variable=backbone.embeddings.patch_embeddings.projection.bias,
+        hf_weight_key="embeddings.patch_embeddings.projection.bias",
+    )
+
+    # Encoder.
+    hidden_dim = backbone.hidden_dim
+    num_heads = backbone.num_heads
+    for i, layer in enumerate(backbone.encoder.layers):
+        prefix = f"encoder.layer.{i}"
+        port_ln(layer.norm1, f"{prefix}.norm1")
+        port_mha(
+            layer.attention.attention,
+            f"{prefix}.attention",
+            num_heads,
+            hidden_dim,
+        )
+        loader.port_weight(
+            keras_variable=layer.layer_scale1.lambda1,
+            hf_weight_key=f"{prefix}.layer_scale1.lambda1",
+        )
+        port_ln(layer.norm2, f"{prefix}.norm2")
+        if backbone.use_swiglu_ffn:
+            port_dense(layer.mlp.weights_in, f"{prefix}.mlp.weights_in")
+            port_dense(layer.mlp.weights_out, f"{prefix}.mlp.weights_out")
+        else:
+            port_dense(layer.mlp.fc1, f"{prefix}.mlp.fc1")
+            port_dense(layer.mlp.fc2, f"{prefix}.mlp.fc2")
+        loader.port_weight(
+            keras_variable=layer.layer_scale2.lambda1,
+            hf_weight_key=f"{prefix}.layer_scale2.lambda1",
+        )
+
+    port_ln(backbone.layernorm, "layernorm")

keras_hub/src/utils/transformers/export/gemma.py

@@ -0,0 +1,89 @@
+import keras.ops as ops
+
+
+def get_gemma_config(backbone):
+    hf_config = {
+        "vocab_size": backbone.vocabulary_size,
+        "num_hidden_layers": backbone.num_layers,
+        "num_attention_heads": backbone.num_query_heads,
+        "num_key_value_heads": backbone.num_key_value_heads,
+        "hidden_size": backbone.hidden_dim,
+        "intermediate_size": backbone.intermediate_dim // 2,
+        "head_dim": backbone.head_dim,
+        "max_position_embeddings": 8192,
+    }
+    return hf_config
+
+
+def get_gemma_weights_map(backbone):
+    weights_dict = {}
+
+    # Map token embedding
+    token_embedding_layer = backbone.get_layer("token_embedding")
+    weights_dict["model.embed_tokens.weight"] = token_embedding_layer.weights[0]
+
+    for i in range(backbone.num_layers):
+        decoder_layer = backbone.get_layer(f"decoder_block_{i}")
+
+        # Pre-attention normalization
+        weights_dict[f"model.layers.{i}.input_layernorm.weight"] = (
+            decoder_layer.pre_attention_norm.weights[0]
+        )
+
+        # Attention query projection
+        query_kernel = decoder_layer.attention.query_dense.weights[0]
+        query_kernel = ops.transpose(query_kernel, axes=(1, 0, 2))
+        query_kernel = ops.reshape(query_kernel, (-1, backbone.hidden_dim))
+        query_kernel = ops.transpose(query_kernel)
+        weights_dict[f"model.layers.{i}.self_attn.q_proj.weight"] = query_kernel
+
+        # Attention key projection
+        key_kernel = decoder_layer.attention.key_dense.weights[0][0]
+        weights_dict[f"model.layers.{i}.self_attn.k_proj.weight"] = (
+            ops.transpose(key_kernel)
+        )
+
+        # Attention value projection
+        value_kernel = decoder_layer.attention.value_dense.weights[0][0]
+        weights_dict[f"model.layers.{i}.self_attn.v_proj.weight"] = (
+            ops.transpose(value_kernel)
+        )
+
+        # Attention output projection
+        out_kernel = decoder_layer.attention.output_dense.weights[0]
+        out_kernel = ops.transpose(out_kernel, axes=(2, 0, 1))
+        out_kernel = ops.reshape(out_kernel, (backbone.hidden_dim, -1))
+        weights_dict[f"model.layers.{i}.self_attn.o_proj.weight"] = out_kernel
+
+        # Post-attention normalization
+        weights_dict[f"model.layers.{i}.post_attention_layernorm.weight"] = (
+            decoder_layer.pre_ffw_norm.weights[0]
+        )
+
+        # MLP gate projection
+        gate_kernel = decoder_layer.gating_ffw.weights[0]
+        weights_dict[f"model.layers.{i}.mlp.gate_proj.weight"] = ops.transpose(
+            gate_kernel
+        )
+
+        # MLP up projection
+        up_kernel = decoder_layer.gating_ffw_2.weights[0]
+        weights_dict[f"model.layers.{i}.mlp.up_proj.weight"] = ops.transpose(
+            up_kernel
+        )
+
+        # MLP down projection
+        down_kernel = decoder_layer.ffw_linear.weights[0]
+        weights_dict[f"model.layers.{i}.mlp.down_proj.weight"] = ops.transpose(
+            down_kernel
+        )
+
+    # Map final normalization
+    weights_dict["model.norm.weight"] = backbone.get_layer(
+        "final_normalization"
+    ).weights[0]
+
+    # Tie weights, but clone to avoid sharing memory issues
+    weights_dict["lm_head.weight"] = ops.copy(token_embedding_layer.weights[0])
+
+    return weights_dict