Add Mistral4

ACKNOWLEDGMENTS.md

@@ -12,7 +12,7 @@ MLX LM was developed with contributions from the following individuals:
 OpenBMB's `MiniCPM` and `MiniCPM3`, Kyutai's `Helium`, State-Space's `Mamba v1` and 
 `Mamba v2`, Z.ai & THUKEG's `GLM`, `GLM4`, `GLM5 (GLM MoE DSA)`, Rednote `dots.llm1`, Baidu's `Ernie4.5 MoE`, 
 inclusionAI's `Bailing MoE e.g. Ling-family`, `Bailing MoE Linear e.g. Ling-Linear-family`, 
-Klear team - Kuaishou Technology's `Klear`, AI21 Lab's `Jamba` IBM's `Granite MoE`, 
+Klear team - Kuaishou Technology's `Klear`, AI21 Lab's `Jamba` IBM's `Granite MoE`, Mistral AI's `Mistral4`, 
 Meituan's `LongCat`, Nvidia's `Nemotron H`, Swiss-AI's `Apertus`, Nikity's `Lille130m`, 
 Alibaba Qwen's `Qwen3Next`, Tele-AI's `TeleChat3`, and Allenai's `OLMoE` and `Olmo 3`; 
 Helped add support for the following model architectures: 

mlx_lm/models/mistral3.py

@@ -1,13 +1,12 @@
-# Copyright © 2025 Apple Inc.
+# Copyright © 2026 Apple Inc.
 
 from dataclasses import dataclass
 from typing import Optional
 
 import mlx.core as mx
 import mlx.nn as nn
-from mlx.utils import tree_flatten, tree_unflatten
 
-from . import llama, ministral3
+from . import llama, ministral3, mistral4
 from .base import BaseModelArgs
 
 
@@ -30,6 +29,10 @@ def __init__(self, args: ModelArgs):
             self.language_model = ministral3.Model(
                 ministral3.ModelArgs.from_dict(args.text_config)
             )
+        elif args.text_config.get("model_type") == "mistral4":
+            self.language_model = mistral4.Model(
+                mistral4.ModelArgs.from_dict(args.text_config)
+            )
         else:
             self.language_model = llama.Model(
                 llama.ModelArgs.from_dict(args.text_config)
@@ -46,12 +49,18 @@ def __call__(
         )
 
     def sanitize(self, weights):
-        weights = tree_unflatten(list(weights.items()))
-        weights.pop("vision_tower", None)
-        weights.pop("multi_modal_projector", None)
-        lm_weights = dict(tree_flatten(weights["language_model"]))
-        weights["language_model"] = self.language_model.sanitize(lm_weights)
-        return dict(tree_flatten(weights))
+        lm_weights = {}
+        for key, value in weights.items():
+            if "vision_tower" in key or "multi_modal_projector" in key:
+                continue
+            if key.startswith("model.language_model."):
+                key = "model." + key.removeprefix("model.language_model.")
+            else:
+                key = key.removeprefix("language_model.")
+            lm_weights[key] = value
+
+        sanitized_lm = self.language_model.sanitize(lm_weights)
+        return {"language_model." + k: v for k, v in sanitized_lm.items()}
 
     @property
     def layers(self):