Unverified Commit f5b2c4c3 authored by kirillrdy's avatar kirillrdy Committed by GitHub
Browse files

python3Packages.mlx: enable on linux (#479305)

parents 95712b41 053f5223

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

+34 −2
Original line number Diff line number Diff line
@@ -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"

    ];

  };

})