pythonPackages.mlx: upgrade and fix build

pkgs/by-name/op/openmpi/package.nix

@@ -26,6 +26,7 @@
   cudaPackages,
   # Enable the Sun Grid Engine bindings
   enableSGE ? false,
+  enablePRRTE ? true,
   # Pass PATH/LD_LIBRARY_PATH to point to current mpirun by default
   enablePrefix ? false,
   # Enable libfabric support (necessary for Omnipath networks) on x86_64 linux
@@ -85,14 +86,14 @@ stdenv.mkDerivation (finalAttrs: {
       zlib
       libevent
       hwloc
+      prrte
     ]
     ++ lib.optionals stdenv.hostPlatform.isLinux [
       libnl
       numactl
       pmix
       ucx
       ucc
-      prrte
     ]
     ++ lib.optionals cudaSupport [ cudaPackages.cuda_cudart ]
     ++ lib.optionals (stdenv.hostPlatform.isLinux || stdenv.hostPlatform.isFreeBSD) [ rdma-core ]
@@ -119,7 +120,7 @@ stdenv.mkDerivation (finalAttrs: {
     "--with-pmix=${lib.getDev pmix}"
     "--with-pmix-libdir=${lib.getLib pmix}/lib"
     # Puts a "default OMPI_PRTERUN" value to mpirun / mpiexec executables
-    (lib.withFeatureAs stdenv.hostPlatform.isLinux "prrte" (lib.getBin prrte))
+    (lib.withFeatureAs enablePRRTE "prrte" (lib.getBin prrte))
     (lib.withFeature enableSGE "sge")
     (lib.enableFeature enablePrefix "mpirun-prefix-by-default")
     # TODO: add UCX support, which is recommended to use with cuda for the most robust OpenMPI build

pkgs/by-name/pr/prrte/package.nix

@@ -75,6 +75,6 @@ stdenv.mkDerivation rec {
     homepage = "https://docs.prrte.org/";
     license = lib.licenses.bsd3;
     maintainers = with lib.maintainers; [ markuskowa ];
-    platforms = lib.platforms.linux;
+    platforms = lib.platforms.unix;
   };
 }

pkgs/development/python-modules/mlx/default.nix

@@ -9,6 +9,9 @@
   blas,
   lapack,
   setuptools,
+  nanobind,
+  openmpi,
+  apple-sdk_14,
 }:
 
 let
@@ -25,29 +28,31 @@ let
     rev = "v3.11.3";
     hash = "sha256-7F0Jon+1oWL7uqet5i1IgHX0fUw/+z0QwEcA3zs5xHg=";
   };
+  fmt = fetchFromGitHub {
+    owner = "fmtlib";
+    repo = "fmt";
+    rev = "10.2.1";
+    hash = "sha256-pEltGLAHLZ3xypD/Ur4dWPWJ9BGVXwqQyKcDWVmC3co=";
+  };
 in
 buildPythonPackage rec {
   pname = "mlx";
-  version = "0.18.0";
+  version = "0.21.1";
 
   src = fetchFromGitHub {
     owner = "ml-explore";
     repo = "mlx";
     rev = "refs/tags/v${version}";
-    hash = "sha256-eFKjCrutqrmhZKzRrLq5nYl0ieqLvoXpbnTxA1NEhWo=";
+    hash = "sha256-wxv9bA9e8VyFv/FMh63sUTTNgkXHGQJNQhLuVynczZA=";
   };
 
   pyproject = true;
 
-  patches = [
-    # With Darwin SDK 11 we cannot include vecLib/cblas_new.h, this needs to wait for PR #229210
-    # In the meantime, pretend Accelerate is not available and use blas/lapack instead.
-    ./disable-accelerate.patch
-  ];
-
   postPatch = ''
     substituteInPlace CMakeLists.txt \
-      --replace "/usr/bin/xcrun" "${xcbuild}/bin/xcrun" \
+      --replace-fail "/usr/bin/xcrun" "${xcbuild}/bin/xcrun"
+    substituteInPlace pyproject.toml \
+      --replace-fail "nanobind==2.2.0" "nanobind"
   '';
 
   dontUseCmakeConfigure = true;
@@ -59,6 +64,7 @@ buildPythonPackage rec {
       (lib.cmakeBool "MLX_BUILD_METAL" false)
       (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_GGUFLIB" "${gguf-tools}")
       (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_JSON" "${nlohmann_json}")
+      (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_FMT" "${fmt}")
     ];
   };
 
@@ -67,14 +73,15 @@ buildPythonPackage rec {
     pybind11
     xcbuild
     zsh
-    gguf-tools
-    nlohmann_json
     setuptools
+    nanobind
+    openmpi
   ];
 
   buildInputs = [
     blas
     lapack
+    apple-sdk_14
   ];
 
   meta = with lib; {
@@ -83,6 +90,6 @@ buildPythonPackage rec {
     changelog = "https://github.com/ml-explore/mlx/releases/tag/v${version}";
     license = licenses.mit;
     platforms = [ "aarch64-darwin" ];
-    maintainers = with maintainers; [ viraptor ];
+    maintainers = with maintainers; [ viraptor Gabriella439 ];
   };
 }