mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-18 21:06:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #25096 from nitins17:update-rocm-ci-scripts

Browse Source 
PiperOrigin-RevId: 700725187
This commit is contained in:
jax authors 2024-11-27 09:27:02 -08:00
commit 04a4f9bd8f
2 changed files with 9 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
build/rocm/dev_build_rocm.py
View File

@ -77,13 +77,14 @@ def build_jax_xla(xla_path, rocm_version, rocm_target, use_clang, clang_path):
build_command = [
"python3",
"./build/build.py",
"--enable_rocm",
"--build_gpu_plugin",
"--gpu_plugin_rocm_version=60",
"build"
f"--use_clang={str(use_clang).lower()}",
"--wheels=jaxlib,jax-rocm-plugin,jax-rocm-pjrt"
"--rocm_path=%/opt/rocm-{rocm_version}/",
"--rocm_version=60",
f"--rocm_amdgpu_targets={rocm_target}",
f"--rocm_path=/opt/rocm-{rocm_version}/",
bazel_options,
"--verbose"
]
if clang_option:

7
build/rocm/tools/build_wheels.py
View File

@ -93,11 +93,12 @@ def build_jaxlib_wheel(
cmd = [
"python",
"build/build.py",
"--enable_rocm",
"--build_gpu_plugin",
"--gpu_plugin_rocm_version=60",
"build"
"--wheels=jaxlib,jax-rocm-plugin,jax-rocm-pjrt"
"--rocm_path=%s" % rocm_path,
"--rocm_version=60",
"--use_clang=%s" % use_clang,
"--verbose"
]
# Add clang path if clang is used.