pyrox.dev
1{
2 lib,
3 stdenv,
4 buildPythonPackage,
5 fetchFromGitHub,
6 replaceVars,
7
8 # build-system
9 setuptools,
10
11 # nativeBuildInputs
12 cmake,
13
14 # buildInputs
15 apple-sdk_14,
16 fmt_11,
17 nanobind,
18 nlohmann_json,
19 pybind11,
20
21 # tests
22 numpy,
23 pytestCheckHook,
24 python,
25 runCommand,
26}:
27
28let
29 # static dependencies included directly during compilation
30 gguf-tools = fetchFromGitHub {
31 owner = "antirez";
32 repo = "gguf-tools";
33 rev = "8fa6eb65236618e28fd7710a0fba565f7faa1848";
34 hash = "sha256-15FvyPOFqTOr5vdWQoPnZz+mYH919++EtghjozDlnSA=";
35 };
36
37 mlx = buildPythonPackage rec {
38 pname = "mlx";
39 version = "0.28.0";
40 pyproject = true;
41
42 src = fetchFromGitHub {
43 owner = "ml-explore";
44 repo = "mlx";
45 tag = "v${version}";
46 hash = "sha256-+2dVZ89a09q8mWIbv6fBsySp7clzRV1tOyqr5hjFrNU=";
47 };
48
49 patches = [
50 (replaceVars ./darwin-build-fixes.patch {
51 sdkVersion = apple-sdk_14.version;
52 })
53 ];
54
55 postPatch = ''
56 substituteInPlace pyproject.toml \
57 --replace-fail "nanobind==2.4.0" "nanobind>=2.4.0"
58
59 substituteInPlace mlx/backend/cpu/jit_compiler.cpp \
60 --replace-fail "g++" "$CXX"
61 '';
62
63 dontUseCmakeConfigure = true;
64
65 enableParallelBuilding = true;
66
67 # Allows multiple cores to be used in Python builds.
68 postUnpack = ''
69 export MAKEFLAGS+="''${enableParallelBuilding:+-j$NIX_BUILD_CORES}"
70 '';
71
72 # updates the wrong fetcher rev attribute
73 passthru.skipBulkUpdate = true;
74
75 env = {
76 PYPI_RELEASE = version;
77 CMAKE_ARGS = toString [
78 # NOTE The `metal` command-line utility used to build the Metal kernels is not open-source.
79 # To build mlx with Metal support in Nix, you'd need to use one of the sandbox escape
80 # hatches which let you interact with a native install of Xcode, such as `composeXcodeWrapper`
81 # or by changing the upstream (e.g., https://github.com/zed-industries/zed/discussions/7016).
82 (lib.cmakeBool "MLX_BUILD_METAL" false)
83 (lib.cmakeBool "USE_SYSTEM_FMT" true)
84 (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_GGUFLIB" "${gguf-tools}")
85 (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_JSON" "${nlohmann_json.src}")
86 ];
87 };
88
89 build-system = [
90 setuptools
91 ];
92
93 nativeBuildInputs = [
94 cmake
95 ];
96
97 buildInputs = [
98 apple-sdk_14
99 fmt_11
100 gguf-tools
101 nanobind
102 nlohmann_json
103 pybind11
104 ];
105
106 pythonImportsCheck = [ "mlx" ];
107
108 # Run the mlx Python test suite.
109 nativeCheckInputs = [
110 numpy
111 pytestCheckHook
112 ];
113
114 enabledTestPaths = [
115 "python/tests/"
116 ];
117
118 disabledTests = [
119 # AssertionError
120 "test_numpy_conv"
121 "test_tensordot"
122 ];
123
124 disabledTestPaths = [
125 # AssertionError
126 "python/tests/test_blas.py"
127 ];
128
129 # Additional testing by executing the example Python scripts supplied with mlx
130 # using the version of the library we've built.
131 passthru.tests = {
132 mlxTest =
133 runCommand "run-mlx-examples"
134 {
135 buildInputs = [ mlx ];
136 nativeBuildInputs = [ python ];
137 }
138 ''
139 cp ${src}/examples/python/logistic_regression.py .
140 ${python.interpreter} logistic_regression.py
141 rm logistic_regression.py
142
143 cp ${src}/examples/python/linear_regression.py .
144 ${python.interpreter} linear_regression.py
145 rm linear_regression.py
146
147 touch $out
148 '';
149 };
150
151 meta = {
152 homepage = "https://github.com/ml-explore/mlx";
153 description = "Array framework for Apple silicon";
154 changelog = "https://github.com/ml-explore/mlx/releases/tag/${src.tag}";
155 license = lib.licenses.mit;
156 platforms = [ "aarch64-darwin" ];
157 maintainers = with lib.maintainers; [
158 viraptor
159 Gabriella439
160 cameronyule
161 ];
162 };
163 };
164in
165mlx