From ccb518434dcf3d0c913faf42a9ee48d3a3db1d8a Mon Sep 17 00:00:00 2001 From: Thanh Binh Date: Mon, 3 Nov 2025 09:50:16 -0600 Subject: [PATCH 1/2] Fix ROCm lib names copy pattern to use wildcard version (cherry picked from commit 577992d59b06af54ff9f697f9be965dcb04cca3e) --- build/build.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/build/build.py b/build/build.py index 70e9bc225eb5..4954ac4a9800 100755 --- a/build/build.py +++ b/build/build.py @@ -722,6 +722,9 @@ async def main(): wheel_dir = wheel.replace("cuda", f"cuda{cuda_major_version}").replace( "-", "_" ) + elif "rocm" in wheel: + # Use wildcard pattern to match any ROCm version (rocm60, rocm7, etc.) + wheel_dir = wheel.replace("rocm", "rocm*").replace("-", "_") else: wheel_dir = wheel @@ -739,7 +742,7 @@ async def main(): wheel_version_suffix += ( f"+{wheel_git_hash}{custom_wheel_version_suffix}" ) - if wheel in ["jax", "jax-cuda-pjrt"]: + if wheel in ["jax", "jax-cuda-pjrt", "jax-rocm-pjrt"]: python_tag = "py" else: python_tag = "cp" From e1b42a8acc307ff8852a6e3956ef33584fb3632e Mon Sep 17 00:00:00 2001 From: Thanh Binh Date: Thu, 13 Nov 2025 16:37:16 -0600 Subject: [PATCH 2/2] Fix ROCm wheel directory naming for editable builds --- build/build.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/build/build.py b/build/build.py index 4954ac4a9800..e47895c5f280 100755 --- a/build/build.py +++ b/build/build.py @@ -723,8 +723,12 @@ async def main(): "-", "_" ) elif "rocm" in wheel: - # Use wildcard pattern to match any ROCm version (rocm60, rocm7, etc.) - wheel_dir = wheel.replace("rocm", "rocm*").replace("-", "_") + if args.editable: + # For editable builds, use the actual ROCm version since directory paths cannot contain wildcards + wheel_dir = wheel.replace("rocm", f"rocm{args.rocm_version}").replace("-", "_") + else: + # For non-editable builds, use wildcard pattern to match any ROCm version in glob patterns + wheel_dir = wheel.replace("rocm", "rocm*").replace("-", "_") else: wheel_dir = wheel