46 Commits

Author SHA1 Message Date
jax authors
140955dce0 Merge pull request #23224 from ROCm:ci_build_script
PiperOrigin-RevId: 668061975
2024-08-27 11:12:05 -07:00
Ayaka
859eacb5a1 Fix mypy error 2024-08-27 16:57:53 +01:00
Ruturaj4
9ce8de5fb0 [ROCm] add build file. 2024-08-23 18:11:48 -05:00
Mathew Odden
5c2ffa893f * Add conditional docker interactive mode
Interactive causes bazel to output more
useful info when running locally.

* Fix issue with rocm el8 repo urls

Work around quirk with rocm version
when it ends with 0

* Fix package name conflict

Ubu22 and higher have a package name conflict
between the debian versions and the AMD provided
versions.

* [ROCm] Use clang env
2024-08-22 10:08:41 -05:00
Ruturaj4
e06be544d4 [ROCm] improve gpu script 2024-08-19 15:13:58 -05:00
Ruturaj4
fd7c52d213 [ROCm] Fix python in rocm ci_build script. 2024-08-15 21:33:51 -05:00
Mathew Odden
fafa03c60f Add missing CPython build deps for pyenv 2024-08-12 15:01:34 -05:00
Mathew Odden
701cda8ebd Fix not finding wheels in bazel output 2024-08-12 15:01:34 -05:00
Mathew Odden
df2d140f51 Fix jenkins notty issue 2024-08-12 15:01:34 -05:00
Mathew Odden
319ebf81c1 Add defaults for ROCm build vars 2024-08-12 15:01:34 -05:00
Mathew Odden
abe44f6d9e Add copyright and license headers to new files 2024-08-12 15:01:34 -05:00
Mathew Odden
a1a0a4ecdd Add support for ROCm development builds
Use get_rocm.py changes in ci_build to pull in
development builds for ROCm.

Specify ROCM_BUILD_JOB and ROCM_BUILD_NUM for
activating the development build path.
2024-08-12 15:01:34 -05:00
Mathew Odden
3175f13c59 Add internal release support to get_rocm.py 2024-08-12 15:01:34 -05:00
Mathew Odden
1e58d76772 [ROCm] Change ROCm builds to manylinux wheels 2024-08-12 15:01:34 -05:00
Yue Sheng
eb571c984a Fix lint in run_single_gpu.py
PiperOrigin-RevId: 658933291
2024-08-02 16:03:03 -07:00
Rahul Batra
7d6fa3c05b [ROCm]: Add support to continue on fail, fix script paths and update Dockerfile to add necessary packages 2024-08-01 17:55:15 -05:00
Peter Hawkins
7f4ef63cd8 Run pyupgrade --py310-plus.
Also apply manual fixes to import sorting and unused imports.
2024-06-26 16:10:18 -04:00
Dan Foreman-Mackey
1c640b25cb Fix pre-commit formatting errors
It looks like #21916 introduced some whitespace issues that are causing
our pre-commits to fail.
2024-06-26 06:38:51 -04:00
Ruturaj4
690dba546c [JAX] change rocm build file for pjrt plugin 2024-06-17 16:49:28 +00:00
zahiqbal
ef7694f26a [ROCM]: Generating pytest html logs from unit-tests. 2024-01-24 15:08:35 +00:00
Rahul Batra
b7a7f0bd80 [ROCm]: Dockerfile updates 2024-01-22 16:08:37 +00:00
Rahul Batra
ef79c19093 [ROCm]: Dockerfile and build script updates
Add hipblaslt in Dockerfile
	Update docker file to default to ROCm5.6
 	CI scripts update to handle multiple ROCm versions
2023-09-08 22:56:59 +00:00
Peter Hawkins
f540ae4338 Fix warning about direct invocation of setup.py during jaxlib build.
The jaxlib wheel build currently uses `python setup.py bdist_wheel` to construct the wheel. Change it to use `python -m build -w` instead.

To avoid Python getting confused between the directory named `build` in the bazel tree and the Python `build` module, move `build_wheel.py` into `jaxlib/tools`.

PiperOrigin-RevId: 548133811
2023-07-14 08:31:16 -07:00
Rahul Batra
80bd361364 [ROCm]: Update README.md 2023-06-22 15:53:04 +00:00
Rahul Batra
88a2c8ca5e [ROCm]: Updates defaults in build script 2023-06-20 20:25:40 +00:00
Rahul Batra
b0e541a730 [ROCm]: Updates for container and build script
-Updated dockerfile.ms
	-Updated build script to switch building against XLA repo
  	-Update CI script
	-Update jaxlib setup.py to add rocm version
2023-06-19 18:13:28 +00:00
Rahul Batra
13e45c8953 [ROCm]: Run pmap test on specific number of GPUs 2023-03-30 18:34:47 +00:00
jax authors
edff87eb07 Merge pull request #13613 from ROCmSoftwarePlatform:rocm_rt_build
PiperOrigin-RevId: 510440289
2023-02-17 08:40:28 -08:00
Chao
0dde7a0fb1
Update Dockerfile.ms
update to ROCm5.4
2023-02-17 14:33:33 +00:00
Rahul Batra
023226e181 [ROCm]: Move dockerfile to ROCm5.4 2023-02-09 20:08:35 +00:00
Chao Chen
a2c9fc02e4 jax-rocm runtime/ci dockerfile multistages 2022-12-12 07:45:12 -08:00
Rohit Santhanam
680096cf04 [ROCm] Fix ROCm dockerfile to remove 2020-resolver opt in.
The 2020-resolver opt in has been removed in pip 22.3 since it has now become the default.
2022-10-19 18:11:04 +00:00
Chao Chen
8c13142ae6 fixed build instructions typo 2022-10-11 07:30:49 -07:00
Jason Furmanek
34f6646050 Add default setting for TENSORFLOW_ROCM_COMMIT 2022-10-07 19:57:53 +00:00
jax authors
1f3048bcf2 Merge pull request #12573 from ROCmSoftwarePlatform:rocm-ci-update
PiperOrigin-RevId: 479369638
2022-10-06 11:33:41 -07:00
Rohit Santhanam
b815ac9d8e [ROCm] Upgrade to ROCm 5.3 and associated enhancements 2022-10-01 04:45:26 -07:00
Jason Furmanek
0783f8982d [ROCM] Add TENSORFLOW_ROCM_COMMIT parameter to ROCM ci build 2022-09-29 12:40:40 +00:00
jax authors
6c47dc51cb Merge pull request #12471 from ROCmSoftwarePlatform:rocm-dockerfile-update
PiperOrigin-RevId: 476387200
2022-09-23 09:16:38 -07:00
Peter Hawkins
ba557d5e1b Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
See https://opensource.google/documentation/reference/releasing/contributions#copyright for more details.

PiperOrigin-RevId: 476167538
2022-09-22 12:27:19 -07:00
Jason Furmanek
9a11b61829 [ROCM] Update Dockerfil.rocm to Ubuntu20 2022-09-22 14:29:30 -04:00
Rohit Santhanam
1b3542427e [ROCm] Update Dockerfile.rocm. 2022-08-09 11:09:10 -07:00
Rohit Santhanam
721602ef59 Upgrade ROCm build docker to ROCm version 5.2. 2022-06-28 21:43:07 -07:00
Jeppe Klitgaard
17de89b16a feat: refactor code using pyupgrade
This PR upgrades legacy Python code to 3.7+ code using pyupgrade:
```sh
pyupgrade --py37-plus --keep-runtime-typing **.py
```

a
2022-05-17 22:14:05 +01:00
Rohit Santhanam
190501bc19 Upgrade ROCm build docker to ROCm version 5.1. 2022-03-31 11:56:38 +00:00
jax authors
cf9a900d78 Merge pull request #9584 from ROCmSoftwarePlatform:rocm_refactor_jaxlib
PiperOrigin-RevId: 432236852
2022-03-03 11:11:02 -08:00
Reza Rahimi
a0d9d81f92 Update JAX to use new math libraries in ROCm-5.0. 2022-03-01 20:02:15 +00:00