Skip to content

Instantly share code, notes, and snippets.

@sck-at-ucy

sck-at-ucy/Package.swift

Created November 21, 2025 20:04
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save sck-at-ucy/0b15e2a7f052623a9d4630a89abed03a to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save sck-at-ucy/0b15e2a7f052623a9d4630a89abed03a to your computer and use it in GitHub Desktop.
Download ZIP
Python MLX → Swift MLX via .mlxfn
Raw
README.md

Python MLX → Swift MLX via .mlxfn

Minimal end-to-end example that:

  1. Defines a tiny model in Python MLX
  2. Exports it as a traced function (.mlxfn)
  3. Imports and runs it from Swift MLX on Apple Silicon

Files in this gist

  • python_mlx_export.py – Python MLX script that exports tiny_mlp.mlxfn
  • Package.swift – SwiftPM package definition using mlx-swift
  • Sources/TinyMLXDemo/main.swift – Swift code that imports and runs the .mlxfn

Requirements

  • Apple Silicon Mac (M-series)

  • Python with mlx installed, e.g.:

    pip install mlx
# or
uv add mlx

  • Xcode (with Command Line Tools)

1. Export the model from Python MLX

Place python_mlx_export.py somewhere convenient and run:

python python_mlx_export.py

This creates:

tiny_mlp.mlxfn

This .mlxfn file is the traced MLX function + parameters to be loaded from Swift.

2. Create the Swift package

From the directory where you want the Swift demo:

mkdir TinyMLXDemo
cd TinyMLXDemo
swift package init --type executable

Then:

  1. Replace the generated Package.swift with the Package.swift from this gist.

  2. Replace Sources/TinyMLXDemo/main.swift with the main.swift from this gist.

  3. Copy the exported function file into the package root:

    cp ../tiny_mlp.mlxfn .

Your tree should look like:

TinyMLXDemo/
  Package.swift
  tiny_mlp.mlxfn
  Sources/
    TinyMLXDemo/
      main.swift

3. Ensure Metal toolchain is available (one-time)

mlx-swift needs the Metal toolchain to build its GPU kernels. If you haven’t used Metal from Xcode on this machine before, run once:

sudo xcode-select -s /Applications/Xcode.app
xcodebuild -runFirstLaunch
sudo xcodebuild -downloadComponent MetalToolchain

4. Build and run the Swift binary

From inside TinyMLXDemo:

xcodebuild build -scheme TinyMLXDemo -destination 'platform=macOS'

If the build succeeds, the executable is in:

~/Library/Developer/Xcode/DerivedData/TinyMLXDemo-*/Build/Products/Debug/TinyMLXDemo

Run it:

~/Library/Developer/Xcode/DerivedData/TinyMLXDemo-*/Build/Products/Debug/TinyMLXDemo

You should see output similar to:

TinyMLP output: [array([[-0.612043, -0.674903]], dtype=float32)]

This output comes from the Python-defined MLX model, exported as tiny_mlp.mlxfn and executed from Swift via importFunction(from:).

Raw
main.swift
import Foundation
import MLX
func runImportedTinyMLP() throws {
// .mlxfn exported from Python (assumed in current working dir)
let url = URL(fileURLWithPath: "tiny_mlp.mlxfn")
// Import traced function
let f = try importFunction(from: url)
// 1D data vector
let xFlat = MLXArray([Float32(1.0), 2.0, 3.0, 4.0])
// Reshape to (1, 4) to match Python example
let x = xFlat.reshaped(1, 4)
// Call imported function (throws)
let y = try f(x)
print("TinyMLP output:", y)
}
do {
try runImportedTinyMLP()
} catch {
print("Error running imported MLX function:", error)
}
Raw
Package.swift
// swift-tools-version: 5.10
import PackageDescription
let package = Package(
name: "TinyMLXDemo",
platforms: [
.macOS(.v14)
],
products: [
.executable(
name: "TinyMLXDemo",
targets: ["TinyMLXDemo"]
),
],
dependencies: [
.package(
url: "https://github.com/ml-explore/mlx-swift.git",
from: "0.29.1"
)
],
targets: [
.executableTarget(
name: "TinyMLXDemo",
dependencies: [
.product(name: "MLX", package: "mlx-swift"),
]
)
]
)
Raw
python_mlx_export.py
import mlx.core as mx
import mlx.nn as nn
class TinyMLP(nn.Module):
def __init__(self, in_dim: int = 4, hidden: int = 8, out_dim: int = 2):
super().__init__()
self.net = nn.Sequential(
nn.Linear(in_dim, hidden),
nn.ReLU(),
nn.Linear(hidden, out_dim),
)
def __call__(self, x):
return self.net(x)
# 1. Create model + some example input
model = TinyMLP()
x_example = mx.random.normal((1, 4))
# 2. Lazily initialize parameters
_ = model(x_example)
# 3. Define a pure function f(x) -> y that closes over the model
def forward(x):
return model(x)
# 4. Export as .mlxfn
mx.export_function(
"tiny_mlp.mlxfn",
forward,
x_example,
shapeless=False, # keep shapes baked in (simpler for demo)
)
print("Exported tiny_mlp.mlxfn")
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment