From ffe67f5b397b24e5c739d69bbd8ac99b5044d9da Mon Sep 17 00:00:00 2001 From: Subodh Dubey Date: Wed, 19 Nov 2025 10:55:05 +0000 Subject: [PATCH 1/2] Jaxlib Rocm Version Extra for the ROCK --- build/build.py | 15 ++++++++++++++- build/tools/utils.py | 11 +++++++++-- 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/build/build.py b/build/build.py index 4954ac4a9800..a14d1f84305c 100755 --- a/build/build.py +++ b/build/build.py @@ -746,16 +746,29 @@ async def main(): python_tag = "py" else: python_tag = "cp" + + logger.info(f"ALL ARGS: {args}") + logger.info(f"Using ROCm path: {args.rocm_path} - {wheel_version_suffix} -{custom_wheel_version_suffix}") + if args.rocm_path: + rocm_tag = args.rocm_path.split("-")[-1] + logger.info(f"Using ROCm tag: {rocm_tag}") + else: + rocm_tag = None + + logger.info(f"Final wheel version suffix: {wheel_version_suffix}") + logger.info(f"Copying wheel files for {wheel_dir} with suffix {wheel_version_suffix} and python tag {python_tag} - Bazel dir: {bazel_dir} - output path: {output_path}") utils.copy_individual_files( bazel_dir, output_path, f"{wheel_dir}*{wheel_version_suffix}-{python_tag}*.whl", - ) + rocm_tag=rocm_tag, + ) if wheel == "jax": utils.copy_individual_files( bazel_dir, output_path, f"{wheel_dir}*{wheel_version_suffix}.tar.gz", + rocm_tag=rocm_tag, ) # Exit with success if all wheels in the list were built successfully. diff --git a/build/tools/utils.py b/build/tools/utils.py index 90654b24cf7b..4005798c4dd7 100644 --- a/build/tools/utils.py +++ b/build/tools/utils.py @@ -265,13 +265,20 @@ def copy_dir_recursively(src, dst): logging.info("Editable wheel path: %s" % dst) -def copy_individual_files(src: str, dst: str, glob_pattern: str): +def copy_individual_files(src: str, dst: str, glob_pattern: str, rocm_tag: str): os.makedirs(dst, exist_ok=True) logging.debug( f"Copying files matching pattern {glob_pattern!r} from {src!r} to {dst!r}" ) for f in glob.glob(os.path.join(src, glob_pattern)): - dst_file = os.path.join(dst, os.path.basename(f)) + logging.debug(f"Found file to copy: {f!r}") + if rocm_tag: + f_list = f.split("-") + f_list[2] = f_list[2]+"+rocm"+rocm_tag + new_f = "-".join(f_list) + else: + new_f = f + dst_file = os.path.join(dst, os.path.basename(new_f)) if os.path.exists(dst_file): os.remove(dst_file) shutil.copy2(f, dst_file) From 11d602a4a48f5572a930067b9bf4dbb78c125407 Mon Sep 17 00:00:00 2001 From: Subodh Dubey Date: Wed, 19 Nov 2025 10:59:13 +0000 Subject: [PATCH 2/2] Jaxlib Rocm Version Extra for the ROCK --- build/build.py | 4 ---- build/tools/utils.py | 2 +- 2 files changed, 1 insertion(+), 5 deletions(-) diff --git a/build/build.py b/build/build.py index a14d1f84305c..17a4d309bb71 100755 --- a/build/build.py +++ b/build/build.py @@ -747,16 +747,12 @@ async def main(): else: python_tag = "cp" - logger.info(f"ALL ARGS: {args}") - logger.info(f"Using ROCm path: {args.rocm_path} - {wheel_version_suffix} -{custom_wheel_version_suffix}") if args.rocm_path: rocm_tag = args.rocm_path.split("-")[-1] logger.info(f"Using ROCm tag: {rocm_tag}") else: rocm_tag = None - logger.info(f"Final wheel version suffix: {wheel_version_suffix}") - logger.info(f"Copying wheel files for {wheel_dir} with suffix {wheel_version_suffix} and python tag {python_tag} - Bazel dir: {bazel_dir} - output path: {output_path}") utils.copy_individual_files( bazel_dir, output_path, diff --git a/build/tools/utils.py b/build/tools/utils.py index 4005798c4dd7..067f19a3df2e 100644 --- a/build/tools/utils.py +++ b/build/tools/utils.py @@ -271,8 +271,8 @@ def copy_individual_files(src: str, dst: str, glob_pattern: str, rocm_tag: str): f"Copying files matching pattern {glob_pattern!r} from {src!r} to {dst!r}" ) for f in glob.glob(os.path.join(src, glob_pattern)): - logging.debug(f"Found file to copy: {f!r}") if rocm_tag: + logging.info(f"Adding Rocm tag {rocm_tag} to file {f}") f_list = f.split("-") f_list[2] = f_list[2]+"+rocm"+rocm_tag new_f = "-".join(f_list)