73 Commits

Author SHA1 Message Date
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