diff --git a/pkgs/development/python-modules/mlx/default.nix b/pkgs/development/python-modules/mlx/default.nix index 95692d73e160..054bd57dcfbd 100644 --- a/pkgs/development/python-modules/mlx/default.nix +++ b/pkgs/development/python-modules/mlx/default.nix @@ -17,6 +17,8 @@ nanobind, nlohmann_json, pybind11, + # linux-only + openblas, # tests numpy, @@ -50,7 +52,7 @@ buildPythonPackage (finalAttrs: { hash = "sha256-Vt0RH+70VBwUjXSfPTsNdRS3g0ookJHhzf2kvgEtgH8="; }; - patches = [ + patches = lib.optionals stdenv.hostPlatform.isDarwin [ (replaceVars ./darwin-build-fixes.patch { sdkVersion = apple-sdk.version; }) @@ -104,6 +106,9 @@ buildPythonPackage (finalAttrs: { nanobind nlohmann_json pybind11 + ] + ++ lib.optionals stdenv.hostPlatform.isLinux [ + openblas ]; pythonImportsCheck = [ "mlx" ]; @@ -118,6 +123,30 @@ buildPythonPackage (finalAttrs: { "python/tests/" ]; + disabledTests = lib.optionals (stdenv.hostPlatform.isLinux && stdenv.hostPlatform.isx86_64) [ + # Segmentation fault + "test_lapack" + "test_multivariate_normal" + "test_orthogonal" + "test_vmap_inverse" + "test_vmap_svd" + ]; + + disabledTestPaths = lib.optionals (stdenv.hostPlatform.isLinux && stdenv.hostPlatform.isx86_64) [ + # Segmentation fault + "python/tests/test_linalg.py" + ]; + + # patchelf is only available on Linux and no patching is needed on darwin. + # Otherwise mlx/core.cpython-313-x86_64-linux-gnu.so contains a reference to + # /build/source/build/temp.linux-x86_64-cpython-313/mlx.core/libmlx.so in its rpath. + postInstall = lib.optionalString stdenv.hostPlatform.isLinux '' + patchelf --replace-needed \ + libmlx.so \ + $out/${python.sitePackages}/mlx/lib64/libmlx.so \ + $out/${python.sitePackages}/mlx/core.cpython-*.so + ''; + # Additional testing by executing the example Python scripts supplied with mlx # using the version of the library we've built. passthru.tests = { @@ -145,12 +174,15 @@ buildPythonPackage (finalAttrs: { description = "Array framework for Apple silicon"; changelog = "https://github.com/ml-explore/mlx/releases/tag/${finalAttrs.src.tag}"; license = lib.licenses.mit; - platforms = [ "aarch64-darwin" ]; maintainers = with lib.maintainers; [ Gabriella439 booxter cameronyule viraptor ]; + badPlatforms = [ + # Building for x86_64 on macOS is not supported + "x86_64-darwin" + ]; }; })