Support building MLX-Swift on Linux through SwiftPM (CPU only)

[Joannis]

Based on #293

Proposed changes

• Linux compilation works.
• Doesn't support cross-compilation yet

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[incertum]

Amazing @Joannis! Let me check in with @davidkoski and ask when we might be ready to bump mlx-c in order to unblock current PRs.

[incertum]

@Joannis you may rebase this PR now since the other PR got merged yesterday. Happy to help with the review once it's rebased.

[Joannis]

Rebased @incertum

[incertum]

Rebased @incertum

wohoo, could you run pre-commit run --all and re-push?

This commit should not exist f85906c after a clean rebase (since you worked on top of the previous PR). You could cherry delete it and or perhaps retry the rebase? Thanks a bunch!

[incertum]

@davidkoski do you have more context around MLXFast issues? Intuitively it sounds like something you would like to have. It seems to just compile fine over CMake on Linux.

[davidkoski]

I am not sure what this is for. MLXFast used to be a separate package, but was folded into MLX, see Sources/MLX/MLXFast.swift.

MLXFast remains as a stub to avoid breaking things and this function is:

@available(*, deprecated, message: "layerNorm is now available in the main MLX module")
@_disfavoredOverload
public func layerNorm(
    _ x: MLXArray, weight: MLXArray? = nil, bias: MLXArray? = nil, eps: Float,
    stream: StreamOrDevice = .default
) -> MLXArray {
    return MLXFast.layerNorm(x, weight: weight, bias: bias, eps: eps, stream: stream)
}

What issue is this working around?

[davidkoski]

In this case MLXFast is an enum to give a namespace to look like the old module MLXFast. So we have:

• MLXFast.layerNorm (the module MLXFast if you were to import that)
• MLX.MLXFast.layerNorm aka MLXFast.layerNorm (the new home, you should not import MLXFast and it is not here)
• MLX.layerNorm (just forwards to MLX.MLXFast.layerNorm for convenience)

OK, so the code here is referencing the second one as there is no import of MLXFast, thus we are getting the MLXFast scoped to the MLX module.

[incertum]

Could we move the Dockerfile somewhere under .github for CI use or for documentation purposes?

How much work would it be to add the swift build to https://github.com/ml-explore/mlx-swift/blob/main/.github/scripts/setup%2Bbuild-linux-container-cmake.sh

or better yet a new script / CI job (my preference)?

[davidkoski]

I think these references to MLXFast are:

1. stale
2. maybe causing the problems in Normalization.swift. I propose the following:

diff --git a/Package.swift b/Package.swift
index d88d9e4..0d1ed11 100644
--- a/Package.swift
+++ b/Package.swift
@@ -188,7 +188,7 @@ let package = Package(
         ),
         .target(
             name: "MLXNN",
-            dependencies: ["MLX", "MLXRandom", "MLXFast"],
+            dependencies: ["MLX"],
             swiftSettings: [
                 .enableExperimentalFeature("StrictConcurrency")
             ]
@@ -218,7 +218,7 @@ let package = Package(
         .testTarget(
             name: "MLXTests",
             dependencies: [
-                "MLX", "MLXRandom", "MLXNN", "MLXOptimizers", "MLXFFT", "MLXLinalg", "MLXFast",
+                "MLX", "MLXNN", "MLXOptimizers",
             ]
         ),
 

[davidkoski]

FYI #319 will have more Package.swift changes -- if we can get this merged first you won't have to figure out how to merge those in.

Outstanding issues. These have to be fixed before the merge:

• swift-format
• figure out this MLXFast issue
• move Dockerfile under .github

Ideally this would be fixed but we can do a followup PR:

• integrate swift build into GitHub actions