DOCSEARCHbyAlgolia
Submit Docs

MLX Swift Examples

repository·main·Indexed 25 days ago

https://github.com/ml-explore/mlx-swift-examples

A collection of demonstration applications and reusable libraries for MLX Swift. It covers implementations for Large Language Models (LLMs), Vision-Language Models (VLMs), Stable Diffusion, and numerical computing, including examples like LLMBasic, LLMEval, LoRATrainingExample, MLX Chat, MNISTTrainer, and curve fitting via gradient descent.

Tokens
14.5K
Snippets
38
Records
76
Agent score
83%

What's inside mlx-swift-examples

  1. Overview of MLX Swift Examples

    main

    MLX Swift Examples provides a collection of demonstration programs using the MLX Swift framework. The repository includes applications for LLMs, VLMs, Stable Diffusion, and general numerical computing tasks.

    Key application categories include:

    • LLM/VLM Applications: LLMBasic, LLMEval, MLXChatExample, and LoRATrainingExample.
    • Image Generation: StableDiffusionExample.
    • Training Examples: MNISTTrainer and LinearModelTraining.
    • Numerical Computing: Examples for CurveFit, HeatTransfer, and Mandelbrot set rendering to demonstrate MLX array idioms, compile, and custom Metal kernels.

  2. Overview of LinearModelTraining tool

    main

    The LinearModelTraining command line tool demonstrates fundamental machine learning concepts using a simple linear model defined by the function f(x) = mx + b. It is designed to illustrate the following core components of ML workflows:

    • Model with parameters: Implementing a simple linear model.
    • Loss function: Measuring the error between predictions and targets.
    • Gradients: Calculating the direction for parameter updates.
    • Optimizers: Using optimization algorithms to minimize loss.
    • Training loop: The iterative process of training the model against an unknown linear function.

  3. Overview of MLXMNIST library

    main

    MLXMNIST is a Swift port of the MNIST training code from the Python MLX examples. It implements a LeNet architecture (instead of a standard MLP) for handwritten digit recognition. The library provides functionality to:

    • Download MNIST test and training datasets.
    • Build the LeNet model architecture.
    • Shuffle and batch data for training and evaluation.

  5. Compare Mandelbrot set implementations in MLX Swift

    main

    The Mandelbrot example provides four different implementation strategies to render the set by iterating z ← z² + c for every pixel in parallel. You can use these to compare performance across different compute backends:

    • Plain MLX (computeMandelbrotMLX): A straightforward implementation using complex64 and linspace to build the grid c, then looping the recurrence over the whole grid.
    • Compiled MLX (computeMandelbrotMLXCompiled): Uses compile(...) to wrap the math. Operations are fused, typically resulting in ~3–4× faster performance on the inner loop compared to plain MLX.
    • Metal kernel (computeMandelbrotMetal): Uses a custom MLXFast.metalKernel. This implementation is optimized by keeping counts in local variables (avoiding per-iteration writes) and allowing pixels to early-exit. It is approximately ~10× faster than the compiled MLX version.
    • Reference CPU (Mandelbrot+CPU.swift): A plain Swift implementation used for correctness verification and comparison.

  6. Understand the Gradient Descent Curve Fitting Algorithm

    main

    The Gradient Descent example demonstrates fitting a quadratic model θ₀ + θ₁·x + θ₂·x² to noisy samples drawn from a cubic target. It uses MLX.grad to automatically compute the gradient of the mean-squared-error loss function without requiring manual derivatives.

    Key components of the algorithm implementation in Algorithm/Gradient.swift:

    • model(_:_:): Defines the quadratic function being fitted.
    • target(_:): Defines the cubic ground truth.
    • Gradient.init(): Samples the target with uniform noise and constructs the gradient function gradLoss = grad(loss).
    • Gradient.step(): Performs a single parameter update using the rule θ ← θ − η · ∇L(θ) and evaluates the new parameters.

  7. Understand MLX Swift Core Concepts

    main

    MLX Swift is a machine learning framework optimized for Apple Silicon. Key architectural concepts include:

    • Lazy Evaluation: Operations are recorded in a computation graph and only executed when results are explicitly requested using .eval().
    • Unified Memory: Leverages Apple Silicon's architecture to avoid costly data transfers between CPU and GPU.
    • Multi-device Support: Supports both .cpu and .gpu devices.
    • Broadcasting: Automatically expands array shapes during arithmetic operations, similar to NumPy.

  10. Interact with the Gradient Descent Visualization UI

    main

    The user interface, implemented in ContentView.swift, uses Swift Charts to visualize the fitting process. The actual noisy samples are plotted in blue, and the model's prediction is plotted in orange.

    Available UI actions:

    • Start: Executes totalSteps updates with a fixed delay between frames to visualize convergence.
    • Reset: Redraws fresh noisy samples and resets the parameters θ to zero.

  11. Generate text from a model

    main

    You can generate text using a synchronous generateText function or an asynchronous streamingGenerate function for real-time UI updates.

    Synchronous Generation

    Use generateText for simple batch processing. It requires a model, a tokenizer, a prompt, and optional parameters for maxTokens, temperature, and topP (nucleus sampling).

    Streaming Generation

    Use streamingGenerate to yield tokens one by one via an onToken callback. This is ideal for chat interfaces. It is an async function and uses Task.yield() to allow the UI to remain responsive.