adopt mlx-swift 0.25.4 (#313)

* adopt mlx-swift 0.25.4

* - fix #274
* - fix #273

* update imports per current mlx-swift structure
* use boolean mask
* add missing model config, debug switch layers with recent updates
* add some simple eval tests -- make sure we can run values through a model, lora train
* linear dimensions now account for quantization

Author: davidkoski

Applications/LLMEval/ContentView.swift

@@ -4,7 +4,6 @@ import AsyncAlgorithms
 import MLX
 import MLXLLM
 import MLXLMCommon
-import MLXRandom
 import MarkdownUI
 import Metal
 import SwiftUI

Applications/LoRATrainingExample/ContentView.swift

@@ -5,7 +5,6 @@ import MLXLLM
 import MLXLMCommon
 import MLXNN
 import MLXOptimizers
-import MLXRandom
 import SwiftUI
 import Tokenizers
 

Applications/MNISTTrainer/ContentView.swift

@@ -4,7 +4,6 @@ import MLX
 import MLXMNIST
 import MLXNN
 import MLXOptimizers
-import MLXRandom
 import SwiftUI
 
 struct TrainingView: View {

Applications/VLMEval/ContentView.swift

@@ -5,7 +5,6 @@ import AsyncAlgorithms
 import CoreImage
 import MLX
 import MLXLMCommon
-import MLXRandom
 import MLXVLM
 import PhotosUI
 import SwiftUI

Libraries/Embedders/Bert.swift

@@ -1,7 +1,6 @@
 // Copyright © 2024 Apple Inc.
 
 import MLX
-import MLXFast
 import MLXNN
 
 extension MLXArray {

Libraries/Embedders/Load.swift

@@ -4,7 +4,6 @@ import Foundation
 @preconcurrency import Hub
 import MLX
 import MLXNN
-import MLXRandom
 import Tokenizers
 
 struct EmbedderError: Error {
@@ -119,26 +118,3 @@ public func loadModelContainer(
     return try await ModelContainer(
         hub: hub, modelDirectory: modelDirectory, configuration: configuration)
 }
-
-// TODO remove once mlx-swift update is adopted
-func quantize(
-    model: Module,
-    filter: (String, Module) -> (groupSize: Int, bits: Int)?,
-    apply: (Module, Int, Int) -> Module? = quantizeSingle(layer:groupSize:bits:)
-) {
-    let updates =
-        model
-        .leafModules()
-        .flattened()
-        .compactMap { (path, m) -> (String, Module)? in
-            if let (groupSize, bits) = filter(path, m) {
-                if let quantized = apply(m, groupSize, bits) {
-                    return (path, quantized)
-                }
-            }
-
-            return nil
-        }
-
-    model.update(modules: ModuleChildren.unflattened(updates))
-}

Libraries/Embedders/NomicBert.swift

@@ -2,7 +2,6 @@
 
 import Foundation
 import MLX
-import MLXFast
 import MLXNN
 
 class NomicEmbedding: Module {

Libraries/Embedders/Pooling.swift

@@ -2,7 +2,6 @@
 
 import Foundation
 import MLX
-import MLXFast
 import MLXLinalg
 import MLXNN
 

Libraries/MLXLLM/LLMModelFactory.swift

@@ -236,6 +236,7 @@ public class LLMRegistry: AbstractModelRegistry, @unchecked Sendable {
             qwen3_1_7b_4bit,
             qwen3_4b_4bit,
             qwen3_8b_4bit,
+            qwen3MoE_30b_a3b_4bit,
             smolLM_135M_4bit,
             mimo_7b_sft_4bit,
             glm4_9b_4bit,

Libraries/MLXLLM/LoraTrain.swift

@@ -5,7 +5,6 @@ import MLX
 import MLXLMCommon
 import MLXNN
 import MLXOptimizers
-import MLXRandom
 import Tokenizers
 
 /// Equivalent to `lora.py/iterate_batches()`.  Used internally by ``LoRATrain``.