nix mlx compilation for better portability

flake.lock

@@ -55,11 +55,28 @@
         "type": "github"
       }
     },
+    "nixpkgs-old": {
+      "locked": {
+        "lastModified": 1767313136,
+        "narHash": "sha256-16KkgfdYqjaeRGBaYsNrhPRRENs0qzkQVUooNHtoy2w=",
+        "owner": "NixOS",
+        "repo": "nixpkgs",
+        "rev": "ac62194c3917d5f474c1a844b6fd6da2db95077d",
+        "type": "github"
+      },
+      "original": {
+        "owner": "NixOS",
+        "ref": "nixos-25.05",
+        "repo": "nixpkgs",
+        "type": "github"
+      }
+    },
     "root": {
       "inputs": {
         "fenix": "fenix",
         "flake-utils": "flake-utils",
         "nixpkgs": "nixpkgs",
+        "nixpkgs-old": "nixpkgs-old",
         "treefmt-nix": "treefmt-nix"
       }
     },

flake.nix

@@ -3,6 +3,7 @@
 
   inputs = {
     nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
+    nixpkgs-old.url = "github:NixOS/nixpkgs/nixos-25.05";
     flake-utils.url = "github:numtide/flake-utils";
     # Provides Rust dev-env integration:
     fenix = {
@@ -36,6 +37,7 @@
     inputs.flake-utils.lib.eachSystem systems (
       system:
       let
+        xcode = pkgs.callPackage ./nix/xcode.nix { };
         pkgs = import inputs.nixpkgs {
           inherit system;
           overlays = [ inputs.fenix.overlays.default ];
@@ -60,6 +62,7 @@
           };
         };
       in
+      rec
       {
         formatter = treefmtEval.config.build.wrapper;
         checks.formatting = treefmtEval.config.build.check inputs.self;
@@ -68,13 +71,18 @@
           ${pkgs.ruff}/bin/ruff check ${inputs.self}/
           touch $out
         '';
+        packages =
+          {
+            mlx = pkgs.python3Packages.callPackage ./nix/mlx.nix { inherit xcode; };
+          };
 
-        devShells.default = pkgs.mkShell {
+        devShells.default = pkgs.mkShellNoCC {
           packages =
             with pkgs;
             [
               # PYTHON
-              python313
+              (python313.withPackages (_: [ packages.mlx ]))
+
               uv
               ruff
               basedpyright
@@ -102,6 +110,7 @@
             ++ (pkgs.lib.optionals pkgs.stdenv.isLinux [
               # IFCONFIG
               unixtools.ifconfig
+              inputs.nixpkgs-old.legacyPackages."x86_64-linux".gcc12
 
               # Build dependencies for Linux
               pkg-config

nix/darwin-build-fixes.patch

@@ -0,0 +1,79 @@
+diff --git a/CMakeLists.txt b/CMakeLists.txt
+index 0ed30932..d8528132 100644
+--- a/CMakeLists.txt
++++ b/CMakeLists.txt
+@@ -177,11 +177,7 @@ if(MLX_BUILD_METAL)
+     add_compile_definitions(MLX_METAL_DEBUG)
+   endif()
+ 
+-  # Throw an error if xcrun not found
+-  execute_process(
+-    COMMAND zsh "-c" "/usr/bin/xcrun -sdk macosx --show-sdk-version"
+-    OUTPUT_VARIABLE MACOS_SDK_VERSION
+-    OUTPUT_STRIP_TRAILING_WHITESPACE COMMAND_ERROR_IS_FATAL ANY)
++  set(MACOS_SDK_VERSION @sdkVersion@)
+ 
+   if(${MACOS_SDK_VERSION} LESS 14.0)
+     message(
+@@ -199,11 +195,8 @@ if(MLX_BUILD_METAL)
+     endif()
+     set(XCRUN_FLAGS "-mmacosx-version-min=${CMAKE_OSX_DEPLOYMENT_TARGET}")
+   endif()
+-  execute_process(
+-    COMMAND
+-      zsh "-c"
+-      "echo \"__METAL_VERSION__\" | xcrun -sdk macosx metal ${XCRUN_FLAGS} -E -x metal -P - | tail -1 | tr -d '\n'"
+-    OUTPUT_VARIABLE MLX_METAL_VERSION COMMAND_ERROR_IS_FATAL ANY)
++  set(
++    MLX_METAL_VERSION @metalVersion@)
+   FetchContent_Declare(metal_cpp URL ${METAL_CPP_URL})
+   FetchContent_MakeAvailable(metal_cpp)
+   target_include_directories(
+diff --git a/cmake/extension.cmake b/cmake/extension.cmake
+index 13db804a..5b385132 100644
+--- a/cmake/extension.cmake
++++ b/cmake/extension.cmake
+@@ -36,7 +36,7 @@ macro(mlx_build_metallib)
+   add_custom_command(
+     OUTPUT ${MTLLIB_BUILD_TARGET}
+     COMMAND
+-      xcrun -sdk macosx metal
++      metal
+       "$<LIST:TRANSFORM,${MTLLIB_INCLUDE_DIRS},PREPEND,-I>"
+       ${MTLLIB_COMPILE_OPTIONS} ${MTLLIB_SOURCES} -o ${MTLLIB_BUILD_TARGET}
+     DEPENDS ${MTLLIB_DEPS} ${MTLLIB_SOURCES}
+diff --git a/mlx/backend/metal/kernels/CMakeLists.txt b/mlx/backend/metal/kernels/CMakeLists.txt
+index 262b0495..5c7446ad 100644
+--- a/mlx/backend/metal/kernels/CMakeLists.txt
++++ b/mlx/backend/metal/kernels/CMakeLists.txt
+@@ -29,7 +29,7 @@ function(build_kernel_base TARGET SRCFILE DEPS)
+                     "-mmacosx-version-min=${CMAKE_OSX_DEPLOYMENT_TARGET}")
+   endif()
+   add_custom_command(
+-    COMMAND xcrun -sdk macosx metal ${METAL_FLAGS} -c ${SRCFILE}
++    COMMAND metal ${METAL_FLAGS} -c ${SRCFILE}
+             -I${PROJECT_SOURCE_DIR} -o ${TARGET}.air
+     DEPENDS ${SRCFILE} ${DEPS} ${BASE_HEADERS}
+     OUTPUT ${TARGET}.air
+@@ -170,7 +170,7 @@ endif()
+ 
+ add_custom_command(
+   OUTPUT ${MLX_METAL_PATH}/mlx.metallib
+-  COMMAND xcrun -sdk macosx metallib ${KERNEL_AIR} -o
++  COMMAND metallib ${KERNEL_AIR} -o
+           ${MLX_METAL_PATH}/mlx.metallib
+   DEPENDS ${KERNEL_AIR}
+   COMMENT "Building mlx.metallib"
+diff --git a/mlx/backend/metal/make_compiled_preamble.sh b/mlx/backend/metal/make_compiled_preamble.sh
+index bb55ed3a..94ea7dd7 100644
+--- a/mlx/backend/metal/make_compiled_preamble.sh
++++ b/mlx/backend/metal/make_compiled_preamble.sh
+@@ -31,7 +31,7 @@ OUTPUT_FILE=${OUTPUT_DIR}/${SRC_NAME}.cpp
+ mkdir -p "$OUTPUT_DIR"
+ 
+ # Use the metal compiler to get a list of headers (with depth)
+-CCC="xcrun -sdk macosx metal -x metal"
++CCC="metal -x metal"
+ HDRS=$( $CCC -I"$SRC_DIR" -I"$JIT_INCLUDES" -DMLX_METAL_JIT -E -P -CC -C -H "$INPUT_FILE" $CFLAGS -w 2>&1 1>/dev/null )
+ 
+ # Remove any included system frameworks (for MetalPerformancePrimitive headers)

nix/mlx.nix

@@ -0,0 +1,184 @@
+{ stdenv
+, lib
+, buildPythonPackage
+, fetchFromGitHub
+, replaceVars
+, fetchzip
+  # build-system
+, setuptools
+  # nativeBuildInputs
+, cmake
+, patchelf
+, zsh
+, # buildInputs
+  nanobind
+, pybind11
+, nlohmann_json
+  # buildInputs (darwin)
+, xcode
+  # buildInputs (linux)
+, openblas
+  # tests
+, numpy
+, pytestCheckHook
+, python
+, runCommand
+, fmt
+}:
+let
+  # static dependencies included directly during compilation
+  gguf-tools = fetchFromGitHub {
+    owner = "antirez";
+    repo = "gguf-tools";
+    rev = "8fa6eb65236618e28fd7710a0fba565f7faa1848";
+    hash = "sha256-15FvyPOFqTOr5vdWQoPnZz+mYH919++EtghjozDlnSA=";
+  };
+
+  metal_cpp = fetchzip {
+    url = "https://developer.apple.com/metal/cpp/files/metal-cpp_26.zip";
+    hash = "sha256-7n2eI2lw/S+Us6l7YPAATKwcIbRRpaQ8VmES7S8ZjY8=";
+
+  };
+
+  macosSdk = "${xcode}/MacOSX.sdk";
+
+  mlx = buildPythonPackage rec {
+    pname = "mlx";
+    version = "0.30.1";
+    pyproject = true;
+
+    src = fetchFromGitHub {
+      owner = "ml-explore";
+      repo = "mlx";
+      tag = "v${version}";
+      hash = "sha256-Vt0RH+70VBwUjXSfPTsNdRS3g0ookJHhzf2kvgEtgH8=";
+    };
+
+    patches = [
+      (replaceVars ./darwin-build-fixes.patch {
+        sdkVersion = if stdenv.isDarwin then xcode.version else "nah";
+        metalVersion = if stdenv.isDarwin then 230 else "nah";
+      })
+    ];
+
+    postPatch = ''
+      substituteInPlace pyproject.toml \
+        --replace-fail "nanobind==2.10.2" "nanobind"
+
+      substituteInPlace mlx/backend/cpu/jit_compiler.cpp \
+        --replace-fail "g++" "$CXX"
+    '';
+
+    preFixup = ''
+      patchelf --set-rpath '$ORIGIN/lib64:$ORIGIN/lib:${lib.makeLibraryPath ([ stdenv.cc.cc.lib ] ++ lib.optionals stdenv.isLinux [ openblas ])}' $out/lib/python*/site-packages/mlx/core*.so
+    '';
+
+    dontUseCmakeConfigure = true;
+
+    enableParallelBuilding = true;
+
+    # Allows multiple cores to be used in Python builds.
+    postUnpack = ''
+      export MAKEFLAGS+="''${enableParallelBuilding:+-j$NIX_BUILD_CORES}"
+    '';
+
+    # updates the wrong fetcher rev attribute
+    passthru.skipBulkUpdate = true;
+
+    env = {
+      DEV_RELEASE = 1;
+      # NOTE The `metal` command-line utility used to build the Metal kernels is not open-source.
+      # this is what the xcode wrapper is for - it patches in the system metal cli
+      CMAKE_ARGS = toString ([
+        (lib.cmakeBool "USE_SYSTEM_FMT" true)
+        (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_GGUFLIB" "${gguf-tools}")
+        (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_JSON" "${nlohmann_json.src}")
+        (lib.cmakeBool "FETCHCONTENT_FULLY_DISCONNECTED" true)
+      ] ++ lib.optionals stdenv.isDarwin [
+        (lib.cmakeBool "MLX_BUILD_METAL" true)
+        (lib.cmakeOptionType "filepath" "METAL_LIB"
+          "${macosSdk}/System/Library/Frameworks/Metal.framework")
+        (lib.cmakeOptionType "filepath" "FOUNDATION_LIB"
+          "${macosSdk}/System/Library/Frameworks/Foundation.framework")
+        (lib.cmakeOptionType "filepath" "QUARTZ_LIB"
+          "${macosSdk}/System/Library/Frameworks/QuartzCore.framework")
+        (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_METAL_CPP" "${metal_cpp}")
+        (lib.cmakeOptionType "string" "CMAKE_OSX_DEPLOYMENT_TARGET" "${xcode.version}")
+        # (lib.cmakeOptionType "filepath" "CMAKE_OSX_SYSROOT" "${macosSdk}")
+        (lib.cmakeOptionType "filepath" "CMAKE_FRAMEWORK_PATH" "${macosSdk}/System/Library/Frameworks")
+        (lib.cmakeOptionType "string" "CMAKE_C_FLAGS"
+           "-Wno-error=elaborated-enum-base -Wno-error=nullability-completeness")
+        (lib.cmakeOptionType "string" "CMAKE_CXX_FLAGS"
+          "-Wno-error=elaborated-enum-base -Wno-error=nullability-completeness")
+      ] ++ lib.optionals stdenv.isLinux [
+
+        (lib.cmakeBool "MLX_BUILD_METAL" false)
+        (lib.cmakeBool "MLX_BUILD_CUDA" true)
+      ]);
+    } // lib.optionalAttrs stdenv.isDarwin {
+      # DEVELOPER_DIR = "${xcode}/Developer";
+      SDKROOT = "${macosSdk}";
+      MACOSX_DEPLOYMENT_TARGET = xcode.version;
+    };
+
+    build-system = [
+      setuptools
+    ];
+
+    nativeBuildInputs = [
+      cmake
+      patchelf
+      zsh
+    ] ++ lib.optionals stdenv.isDarwin [ xcode ];
+
+    buildInputs = [
+      fmt
+      gguf-tools
+      nanobind
+      pybind11
+    ] ++ lib.optionals stdenv.isLinux [ openblas ];
+
+    pythonImportsCheck = [ "mlx" ];
+
+    # Run the mlx Python test suite.
+    nativeCheckInputs = [
+      numpy
+      pytestCheckHook
+    ];
+
+    enabledTestPaths = [
+      "python/tests/"
+    ];
+
+    # Additional testing by executing the example Python scripts supplied with mlx
+    # using the version of the library we've built.
+    passthru.tests = {
+      mlxTest =
+        runCommand "run-mlx-examples"
+          {
+            buildInputs = [ mlx ];
+            nativeBuildInputs = [ python ];
+          }
+          ''
+            cp ${src}/examples/python/logistic_regression.py .
+            ${python.interpreter} logistic_regression.py
+            rm logistic_regression.py
+
+            cp ${src}/examples/python/linear_regression.py .
+            ${python.interpreter} linear_regression.py
+            rm linear_regression.py
+
+            touch $out
+          '';
+    };
+
+    meta = {
+      homepage = "https://github.com/ml-explore/mlx";
+      description = "Array framework for Apple silicon";
+      changelog = "https://github.com/ml-explore/mlx/releases/tag/${src.tag}";
+      license = lib.licenses.mit;
+      platforms = [ "x86_64-linux" "aarch64-linux" "aarch64-darwin" ];
+    };
+  };
+in
+mlx

nix/xcode.nix

@@ -0,0 +1,22 @@
+{ stdenv
+, xcodeBaseDir ? "/Applications/Xcode.app"
+, xcodeVersion ? "26.2"
+}:
+assert stdenv.hostPlatform.isDarwin;
+stdenv.mkDerivation {
+  pname = "xcode-wrapper-impure";
+  version = xcodeVersion;
+
+  __noChroot = true;
+  buildCommand = ''
+    DEVELOPER_DIR=${xcodeBaseDir}/Contents/Developer/
+    SDKROOT=${xcodeBaseDir}/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX.sdk
+    export DEVELOPER_DIR SDKROOT
+    mkdir -p $out/bin && cd $out/bin
+    ln -s $(/usr/bin/xcrun -f metal) $out/bin/metal
+    ln -s $(/usr/bin/xcrun -f metallib) $out/bin/metallib
+    cd ..
+    ln -s ${xcodeBaseDir}/Contents/Developer/Platforms/MacOSX.platform/Developer/SDKs/MacOSX.sdk $out/MacOSX.sdk
+    ln -s ${xcodeBaseDir}/Contents/Developer/ $out/Developer
+  '';
+}