MLX Swift brings model experiments to a compiled language

Running a model and changing how it learns require different interfaces. A prediction API can hide most of the computation behind a loaded model. Research needs access to the arrays and functions that define that computation, including the derivatives used to update parameters. MLX Swift brings those operations into Swift with an interface intended for experimentation on Apple silicon.

The February 20 announcement describes Swift bindings for MLX, an array framework for machine learning research. It includes neural network and optimizer packages alongside the core array API. A text-generation example demonstrates Mistral 7B, while a separate example trains a network on MNIST. Training is part of the initial offering, so interpreting this release solely as a way to embed a finished model misses much of its purpose.

I like Swift's explicit data model and generally find the language well designed. An experimental numerical API is interesting because it applies that language to a different kind of work from ordinary application interfaces. The relevant question is whether researchers can express computations and their derivatives without leaving Swift for every change to a model.

The computation remains accessible

MLX operates on N-dimensional arrays. The Swift announcement demonstrates creating an integer array with a specified shape, indexing it and adding a randomly generated array. The default device in its example is the GPU. CPU execution is also part of the framework's hardware acceleration model.

Shape is an important part of what the array interface expresses. Six integer values can form three rows of two columns. A second array with two entries can participate in the addition through broadcasting. The operation applies the smaller array across the rows, so the programmer can express the intended calculation without spelling out a separate scalar addition for every element.

The introductory example translates into a compact Swift fragment. This uses the construction and indexing shown in the announcement, without relying on the particular random values printed there.

import MLX

let values = MLXArray(0 ..< 6, [3, 2])
print(values[0, -1])
print(values[0 ..< 2])

Here the negative column index selects the last entry of the first row. The range selects the first two rows. These operations matter beyond convenient syntax because researchers repeatedly inspect and transform parts of larger tensors. Having array operations directly available lets an experiment change the computation at that level.

The announcement does not claim that a Swift compiler can statically prove every array shape is compatible. The example establishes the numerical operations available through MLXArray. It is useful to distinguish the language's type checking from the numerical constraints of a particular tensor calculation. Adopting Swift does not remove the need to understand the shapes a model expects.

Differentiation transforms a function

The more consequential example defines a function that squares an array and passes that function to grad. The result is another callable function, representing its gradient. Evaluating it at a scalar value of 1.5 yields 3. The programmer writes the original computation and asks MLX to derive the corresponding gradient computation.

import MLX

func squared(_ x: MLXArray) -> MLXArray {
    x.square()
}

let derivative = grad(squared)
print(derivative(MLXArray(1.5)))

This is the announcement's differentiation example with descriptive local names. Its simplicity makes the interface clear. grad accepts a function operating on MLX arrays, and the transformed function accepts an MLX array as well. The original function remains ordinary source code that a researcher can change.

Automatic differentiation is directly relevant to gradient-based training. The release pairs it with higher-level neural network and optimizer packages, so the offering includes tools for defining a network and updating it. Those packages avoid making every experiment start from raw array expressions, while preserving access to the lower-level computations.

Function transformations also appear in the announcement's discussion of optimizing computation graphs for speed or memory use. That does not establish a performance result for an arbitrary model. It describes a programming interface through which the framework can operate on computations. The distinction prevents the compiled language argument from becoming an unsupported claim that every Swift experiment will outperform an equivalent program elsewhere.

Swift calls into an existing numerical engine

The implementation includes a C API between Swift and the MLX C++ core. That division explains how a new language interface can expose the framework's numerical capabilities without rewriting the engine in Swift. Researchers use Swift APIs while the implementation reaches the existing core through the bridge.

The repository README at the announcement date describes adding the package to an Xcode project and linking the needed libraries. It separates the core MLX product from MLXNN, MLXOptimizers and MLXRandom. That organization matches the distinction between array calculations and the higher-level machinery used by training code.

I remain wary of treating a pleasant API as evidence that the surrounding development workflow will be pleasant too. Build time is already a source of friction in Swift. The C++ core and GPU work make installation and build requirements part of evaluating this research environment, alongside the syntax of the model itself.

The release places the packages under the MIT license. That is a concrete condition for people who want to inspect or extend the implementation. It does not change the stated deployment scope. Both the announcement and the contemporary README explicitly describe MLX as intended for research, rather than production deployment of models in applications.

The examples define the intended use

The text-generation example supports Mistral and Llama-style models, including pre-quantized MLX models. This makes it possible to explore an existing language model through Swift. It does not establish support for every model architecture or every format a researcher might encounter.

The MNIST example trains a multilayer perceptron to classify handwritten digits. Its role is different from text generation. It shows the neural network and optimizer packages participating in a training computation, which is the part an inference-only example cannot demonstrate.

Together, those examples give prospective users two concrete entry points. Someone examining language-model generation can start with the supplied generation implementation. Someone learning the training API can inspect a smaller classification problem. In both cases, the release makes Swift source available at the level where experiments change the computation, while retaining MLX's accelerated numerical implementation underneath it.