v4.3 tag release update. (#2789)

examples/python/CuTeDSL/cute/ffi/CMakeLists.txt

@@ -35,20 +35,26 @@ find_package(Python3 COMPONENTS Interpreter Development REQUIRED)
 
 # Get Python site-packages directory using Python
 execute_process(
-    COMMAND ${Python3_EXECUTABLE} -c "import site; print(site.getsitepackages()[0])"
-    OUTPUT_VARIABLE Python_SITE_PACKAGES
+    COMMAND ${Python3_EXECUTABLE} -c "import sys, sysconfig; print(';'.join([sysconfig.get_paths()['purelib'],sysconfig.get_paths(vars={'base': sys.base_prefix, 'platbase': sys.base_prefix})['purelib']]))"
+    OUTPUT_VARIABLE Python_SITE_PACKAGES_PATHS
     OUTPUT_STRIP_TRAILING_WHITESPACE
 )
 
-message(STATUS "Python site-packages directory: ${Python_SITE_PACKAGES}")
+message(STATUS "Python site-packages directories: ${Python_SITE_PACKAGES_PATHS}")
 
 # Add nanobind path to CMAKE_PREFIX_PATH
-list(APPEND CMAKE_PREFIX_PATH ${Python_SITE_PACKAGES}/nanobind/cmake)
+foreach(path IN LISTS Python_SITE_PACKAGES_PATHS)
+    if(EXISTS "${path}/nanobind/cmake")
+        message(STATUS "Adding nanobind cmake path: ${path}/nanobind/cmake")
+        list(APPEND CMAKE_PREFIX_PATH "${path}/nanobind/cmake")
+        break()
+    endif()
+endforeach()
 
 # Find nanobind
 find_package(nanobind)
 if(NOT nanobind_FOUND)
-    message(FATAL_ERROR 
+    message(FATAL_ERROR
         "nanobind not found!\n"
         "Please install nanobind with: pip install nanobind\n"
     )

examples/python/CuTeDSL/cute/ffi/jit_argument.py

@@ -243,12 +243,7 @@ import tempfile
 import torch
 
 
-def run_test(tmpdir=None, cmake_args=""):
-    # Skip cleanup if user provides tmpdir
-    cleanup = tmpdir is None
-    # Initialize temporary build directory
-    tmpdir = tmpdir or tempfile.mkdtemp()
-
+def run_test(tmpdir=None, cmake_args="", cleanup=True):
     try:
         current_dir = os.path.dirname(os.path.abspath(__file__))
 
@@ -256,8 +251,6 @@ def run_test(tmpdir=None, cmake_args=""):
         subprocess.run(["cmake", "-B", tmpdir, current_dir] + cmake_args, check=True)
         subprocess.run(["cmake", "--build", tmpdir], check=True)
 
-        sys.path.append(tmpdir)
-
         from tensor import make_tensor, pycapsule_get_pointer
 
         # Mock test tensor and corresponding C structure for this example
@@ -314,4 +307,13 @@ if __name__ == "__main__":
     )
     args = parser.parse_args()
 
-    run_test(tmpdir=args.tmp_dir, cmake_args=args.cmake_args)
+    if args.tmp_dir:
+        tmp_dir = args.tmp_dir
+        cleanup = False
+    else:
+        tmp_dir = tempfile.mkdtemp()
+        cleanup = True
+
+    sys.path.append(tmp_dir)
+
+    run_test(tmp_dir, args.cmake_args, cleanup)