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

{
  lib,
  stdenv,
  buildPythonPackage,
  fetchFromGitHub,
  replaceVars,

  # build-system
  setuptools,

  # nativeBuildInputs
  cmake,

  # buildInputs
  apple-sdk,
  fmt,
  nanobind,
  nlohmann_json,
  pybind11,
  # linux-only
  openblas,

  # tests
  numpy,
  pytestCheckHook,
  python,
  runCommand,

  # passthru
  mlx,
}:

let
  # static dependencies included directly during compilation
  gguf-tools = fetchFromGitHub {
    owner = "antirez";
    repo = "gguf-tools";
    rev = "8fa6eb65236618e28fd7710a0fba565f7faa1848";
    hash = "sha256-15FvyPOFqTOr5vdWQoPnZz+mYH919++EtghjozDlnSA=";
  };

in
buildPythonPackage (finalAttrs: {
  pname = "mlx";
  version = "0.30.1";
  pyproject = true;

  src = fetchFromGitHub {
    owner = "ml-explore";
    repo = "mlx";
    tag = "v${finalAttrs.version}";
    hash = "sha256-Vt0RH+70VBwUjXSfPTsNdRS3g0ookJHhzf2kvgEtgH8=";
  };

  patches = lib.optionals stdenv.hostPlatform.isDarwin [
    (replaceVars ./darwin-build-fixes.patch {
      sdkVersion = apple-sdk.version;
    })
  ];

  postPatch = ''
    substituteInPlace pyproject.toml \
      --replace-fail "nanobind==2.10.2" "nanobind"

    substituteInPlace mlx/backend/cpu/jit_compiler.cpp \
      --replace-fail "g++" "${lib.getExe' stdenv.cc "c++"}"
  '';

  dontUseCmakeConfigure = true;

  enableParallelBuilding = true;

  # Allows multiple cores to be used in Python builds.
  postUnpack = ''
    export MAKEFLAGS+="''${enableParallelBuilding:+-j$NIX_BUILD_CORES}"
  '';

  # updates the wrong fetcher rev attribute
  passthru.skipBulkUpdate = true;

  env = {
    DEV_RELEASE = 1;
    CMAKE_ARGS = toString [
      # NOTE The `metal` command-line utility used to build the Metal kernels is not open-source.
      # To build mlx with Metal support in Nix, you'd need to use one of the sandbox escape
      # hatches which let you interact with a native install of Xcode, such as `composeXcodeWrapper`
      # or by changing the upstream (e.g., https://github.com/zed-industries/zed/discussions/7016).
      (lib.cmakeBool "MLX_BUILD_METAL" false)
      (lib.cmakeBool "USE_SYSTEM_FMT" true)
      (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_GGUFLIB" "${gguf-tools}")
      (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_JSON" "${nlohmann_json.src}")
    ];
  };

  build-system = [
    setuptools
  ];

  nativeBuildInputs = [
    cmake
  ];

  buildInputs = [
    fmt
    gguf-tools
    nanobind
    nlohmann_json
    pybind11
  ]
  ++ lib.optionals stdenv.hostPlatform.isLinux [
    openblas
  ];

  pythonImportsCheck = [ "mlx" ];

  # Run the mlx Python test suite.
  nativeCheckInputs = [
    numpy
    pytestCheckHook
  ];

  enabledTestPaths = [
    "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 = {
    mlxTest =
      runCommand "run-mlx-examples"
        {
          buildInputs = [ mlx ];
          nativeBuildInputs = [ python ];
        }
        ''
          cp ${finalAttrs.src}/examples/python/logistic_regression.py .
          ${python.interpreter} logistic_regression.py
          rm logistic_regression.py

          cp ${finalAttrs.src}/examples/python/linear_regression.py .
          ${python.interpreter} linear_regression.py
          rm linear_regression.py

          touch $out
        '';
  };

  meta = {
    homepage = "https://github.com/ml-explore/mlx";
    description = "Array framework for Apple silicon";
    changelog = "https://github.com/ml-explore/mlx/releases/tag/${finalAttrs.src.tag}";
    license = lib.licenses.mit;
    maintainers = with lib.maintainers; [
      Gabriella439
      booxter
      cameronyule
      viraptor
    ];
    badPlatforms = [
      # Building for x86_64 on macOS is not supported
      "x86_64-darwin"
    ];
  };
})