From fc19a08caa7d706a79baa0945824e89c2688c774 Mon Sep 17 00:00:00 2001
From: Cheng <zcbenz@gmail.com>
Date: Fri, 19 Dec 2025 16:43:00 +0900
Subject: [PATCH] Set install rpath of python bindings with cmake (#2934)

---
 .github/actions/build-cuda-release/action.yml |  2 +-
 .../actions/build-linux-release/action.yml    |  7 +++-
 python/scripts/repair_linux.sh                | 20 -----------
 python/scripts/repair_record.py               | 33 -------------------
 python/src/CMakeLists.txt                     |  6 ++--
 5 files changed, 11 insertions(+), 57 deletions(-)
 delete mode 100644 python/scripts/repair_linux.sh
 delete mode 100644 python/scripts/repair_record.py

diff --git a/.github/actions/build-cuda-release/action.yml b/.github/actions/build-cuda-release/action.yml
index 9f755881..0ed7719d 100644
--- a/.github/actions/build-cuda-release/action.yml
+++ b/.github/actions/build-cuda-release/action.yml
@@ -22,7 +22,7 @@ runs:
         python setup.py clean --all
         MLX_BUILD_STAGE=2 python -m build -w
 
-        auditwheel repair dist/* \
+        auditwheel repair dist/mlx_cuda*.whl \
           --plat manylinux_2_35_${{ inputs.arch }} \
           --exclude libcublas* \
           --exclude libcuda* \
diff --git a/.github/actions/build-linux-release/action.yml b/.github/actions/build-linux-release/action.yml
index 3ea0e6fc..85a0125e 100644
--- a/.github/actions/build-linux-release/action.yml
+++ b/.github/actions/build-linux-release/action.yml
@@ -24,13 +24,18 @@ runs:
         pip install -e ".[dev]" -v
         pip install typing_extensions
         python setup.py generate_stubs
+
     - name: Build Python package
       shell: bash
       run: |
         pip install auditwheel patchelf build
         python setup.py clean --all
         MLX_BUILD_STAGE=1 python -m build -w
-        bash python/scripts/repair_linux.sh ${{ inputs.arch }}
+        auditwheel repair dist/mlx-*.whl \
+          --plat manylinux_2_35_${{ inputs.arch }} \
+          --exclude libmlx.so* \
+          --only-plat
+
     - name: Build backend package
       if: ${{ inputs.build-backend }}
       shell: bash
diff --git a/python/scripts/repair_linux.sh b/python/scripts/repair_linux.sh
deleted file mode 100644
index bbf4dc02..00000000
--- a/python/scripts/repair_linux.sh
+++ /dev/null
@@ -1,20 +0,0 @@
-#!/bin/bash
-
-auditwheel repair dist/* \
-  --plat manylinux_2_35_${1} \
-  --only-plat \
-  --exclude libmlx* \
-  -w wheel_tmp
-
-mkdir wheelhouse
-cd wheel_tmp
-repaired_wheel=$(find . -name "*.whl" -print -quit)
-unzip -q "${repaired_wheel}"
-rm "${repaired_wheel}"
-core_so=$(find mlx -name "core*.so" -print -quit)
-rpath="\$ORIGIN/lib"
-patchelf --force-rpath --set-rpath "$rpath" "$core_so"
-python ../python/scripts/repair_record.py ${core_so}
-
-# Re-zip the repaired wheel
-zip -r -q "../wheelhouse/${repaired_wheel}" .
diff --git a/python/scripts/repair_record.py b/python/scripts/repair_record.py
deleted file mode 100644
index 1738fd5a..00000000
--- a/python/scripts/repair_record.py
+++ /dev/null
@@ -1,33 +0,0 @@
-import base64
-import glob
-import hashlib
-import sys
-
-filename = sys.argv[1]
-
-
-# Compute the new hash and size
-def urlsafe_b64encode(data: bytes) -> bytes:
-    return base64.urlsafe_b64encode(data).rstrip(b"=")
-
-
-hasher = hashlib.sha256()
-with open(filename, "rb") as f:
-    data = f.read()
-    hasher.update(data)
-hash_str = urlsafe_b64encode(hasher.digest()).decode("ascii")
-size = len(data)
-
-# Update the record file
-record_file = glob.glob("*/RECORD")[0]
-with open(record_file, "r") as f:
-    lines = [l.split(",") for l in f.readlines()]
-
-for l in lines:
-    if filename == l[0]:
-        l[1] = hash_str
-        l[2] = f"{size}\n"
-
-with open(record_file, "w") as f:
-    for l in lines:
-        f.write(",".join(l))
diff --git a/python/src/CMakeLists.txt b/python/src/CMakeLists.txt
index 0580b878..7053e61c 100644
--- a/python/src/CMakeLists.txt
+++ b/python/src/CMakeLists.txt
@@ -55,8 +55,10 @@ target_link_libraries(core PRIVATE mlx)
 
 if(BUILD_SHARED_LIBS)
   if(${CMAKE_SYSTEM_NAME} MATCHES "Darwin")
-    target_link_options(core PRIVATE -Wl,-rpath,@loader_path/lib)
+    set_target_properties(core PROPERTIES INSTALL_RPATH "@loader_path/lib")
   else()
-    target_link_options(core PRIVATE -Wl,-rpath,\$ORIGIN/lib)
+    set_target_properties(core PROPERTIES INSTALL_RPATH "\$ORIGIN/lib")
   endif()
+  # Do not add build dir to rpath.
+  set_target_properties(core PROPERTIES BUILD_WITH_INSTALL_RPATH ON)
 endif()