mlx-cpu 0.30.1__py3-none-manylinux_2_35_x86_64.whl

mlx/__main__.py

@@ -0,0 +1,27 @@
+import argparse
+
+
+def main() -> None:
+    from mlx.core import __version__
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--version",
+        action="version",
+        version=__version__,
+        help="Print the version number.",
+    )
+    parser.add_argument(
+        "--cmake-dir",
+        action="store_true",
+        help="Print the path to the MLX CMake module directory.",
+    )
+    args = parser.parse_args()
+    if args.cmake_dir:
+        from pathlib import Path
+
+        print(Path(__file__).parent)
+
+
+if __name__ == "__main__":
+    main()

mlx/_reprlib_fix.py

@@ -0,0 +1,16 @@
+# Copyright © 2023 Apple Inc.
+
+import array
+import reprlib
+
+_old_repr_array = reprlib.Repr.repr_array
+
+
+def repr_array(self, x, maxlevel):
+    if isinstance(x, array.array):
+        return _old_repr_array(self, x, maxlevel)
+    else:
+        return self.repr_instance(x, maxlevel)
+
+
+reprlib.Repr.repr_array = repr_array

mlx/extension.py

@@ -0,0 +1,88 @@
+# Copyright © 2023 Apple Inc.
+
+import os
+import re
+import subprocess
+import sys
+from pathlib import Path
+
+from setuptools import Extension
+from setuptools.command.build_ext import build_ext
+
+import mlx
+
+_MLX_PATH = str(mlx.__path__[0])
+
+
+# A CMakeExtension needs a sourcedir instead of a file list.
+class CMakeExtension(Extension):
+    def __init__(self, name: str, sourcedir: str = "") -> None:
+        super().__init__(name, sources=[])
+        self.sourcedir = os.fspath(Path(sourcedir).resolve())
+
+
+class CMakeBuild(build_ext):
+    def build_extension(self, ext: CMakeExtension) -> None:
+        # Must be in this form due to bug in .resolve() only fixed in Python 3.10+
+        ext_fullpath = Path.cwd() / self.get_ext_fullpath(ext.name)  # type: ignore[no-untyped-call]
+        extdir = ext_fullpath.parent.resolve()
+
+        debug = int(os.environ.get("DEBUG", 0)) if self.debug is None else self.debug
+        cfg = "Debug" if debug else "Release"
+
+        # Set Python_EXECUTABLE instead if you use PYBIND11_FINDPYTHON
+        # EXAMPLE_VERSION_INFO shows you how to pass a value into the C++ code
+        # from Python.
+        cmake_args = [
+            f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={extdir}{os.sep}",
+            f"-DCMAKE_BUILD_TYPE={cfg}",
+            "-DBUILD_SHARED_LIBS=ON",
+        ]
+        build_args = []
+        # Adding CMake arguments set as environment variable
+        # (needed e.g. to build for ARM OSx on conda-forge)
+        if "CMAKE_ARGS" in os.environ:
+            cmake_args += [item for item in os.environ["CMAKE_ARGS"].split(" ") if item]
+
+        if sys.platform.startswith("darwin"):
+            # Cross-compile support for macOS - respect ARCHFLAGS if set
+            archs = re.findall(r"-arch (\S+)", os.environ.get("ARCHFLAGS", ""))
+            if archs:
+                cmake_args += ["-DCMAKE_OSX_ARCHITECTURES={}".format(";".join(archs))]
+
+        # Set CMAKE_BUILD_PARALLEL_LEVEL to control the parallel build level
+        # across all generators.
+        if "CMAKE_BUILD_PARALLEL_LEVEL" not in os.environ:
+            build_args += [f"-j{os.cpu_count()}"]
+
+        build_temp = Path(self.build_temp) / ext.name
+        if not build_temp.exists():
+            build_temp.mkdir(parents=True)
+
+        # Make sure cmake can find MLX
+        os.environ["MLX_DIR"] = _MLX_PATH
+
+        subprocess.run(
+            ["cmake", ext.sourcedir, *cmake_args], cwd=build_temp, check=True
+        )
+        subprocess.run(
+            ["cmake", "--build", ".", *build_args], cwd=build_temp, check=True
+        )
+
+    def run(self) -> None:
+        super().run()
+
+        # Based on https://github.com/pypa/setuptools/blob/main/setuptools/command/build_ext.py#L102
+        if self.inplace:
+            for ext in self.extensions:
+                if isinstance(ext, CMakeExtension):
+                    # Resolve inplace package dir
+                    build_py = self.get_finalized_command("build_py")
+                    inplace_file, regular_file = self._get_inplace_equivalent(
+                        build_py, ext
+                    )
+
+                    inplace_dir = str(Path(inplace_file).parent.resolve())
+                    regular_dir = str(Path(regular_file).parent.resolve())
+
+                    self.copy_tree(regular_dir, inplace_dir)