mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 11:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
25,054 Commits 501 Branches 263 Tags
Commit Graph

68 Commits

Author SHA1 Message Date
Ruturaj4
435edf1f8c Add gfx12xx archs 2025-01-15 16:14:40 +00:00
Ruturaj4
20b75ab82f Update package indentation fix 2025-01-01 18:50:47 -06:00
jax authors
bad6f0f0e9 Merge pull request #25440 from ROCm:jax_pypi-upstream 
PiperOrigin-RevId: 707934347
 2024-12-19 09:17:46 -08:00
Ruturaj4
fefebebea4 Update documentation and add setup.py pypi bindings 2024-12-19 10:21:28 -06:00
Ruturaj4
bfcace4933 [ROCm] ci build and dockerfile changes 2024-12-19 09:02:55 -06:00
jax authors
7b32d88247 Merge pull request #25136 from ROCm:ci_dockerfile_arg_changes-upstream 
PiperOrigin-RevId: 701959495
 2024-12-02 07:21:15 -08:00
jax authors
ab79066bbe Merge pull request #25128 from ROCm:ci_fix_wheelhouse_relative_paths-upstream 
PiperOrigin-RevId: 701143534
 2024-11-28 18:51:49 -08:00
jax authors
b0df405250 Merge pull request #25130 from ROCm:ci_fix_set_options-upstream 
PiperOrigin-RevId: 701143242
 2024-11-28 18:49:47 -08:00
jax authors
385e2f4339 Merge pull request #25137 from ROCm:ci_enable_https-upstream 
PiperOrigin-RevId: 701142951
 2024-11-28 18:45:52 -08:00
Ruturaj4
3d8063209e Update http to https in amd artifactory url. 2024-11-26 22:38:04 +00:00
Ruturaj4
8df2766466 Add argument to override base docker in dockerfile 2024-11-26 22:28:53 +00:00
Ruturaj4
694de6b64c [ROCm] Change run_multi_gpu set opts 2024-11-26 21:10:37 +00:00
Ruturaj4
d30ec2b5b3 [ROCm] fix jax and wheelhouse relative paths 2024-11-26 21:04:36 +00:00
Nitin Srinivasan
84dc9bab33 Update ROCm scripts to match new build.py usage 2024-11-25 19:25:08 +00:00
Ruturaj4
bfd7075c39 [ROCm] ci build fixes 2024-10-25 05:01:44 -05:00
jax authors
3bdc57dd29 Merge pull request #24300 from ROCm:ci_rocm_readme 
PiperOrigin-RevId: 686872994
 2024-10-17 05:21:13 -07:00
Ruturaj4
3c3b08dfd6 [ROCm] Fix README.md to update AMD JAX installation instructions 2024-10-16 17:15:32 -05:00
Ruturaj4
a2824862f5 [ROCm] build script fix 2024-10-15 13:43:08 -05:00
Ruturaj4
937d79e3f2 [ROCm] apt update 2024-10-11 18:33:12 -05:00
Ruturaj4
33bcd0cb7a [ROCm] Bring up clang support for JAX+XLA 
* Add clang path

* bazelrc env fixes

* Fix wheelhouse installation and preserve wheels

* dockerfile changes

* Add target.lst

* Change target architectures

* Install bzip2 and sqlite packages
 2024-10-10 16:31:26 -05:00
Mathew Odden
9ff891dfa1 [ROCm] Remove broken legacy env vars 
These env vars are no longer used or need and were
being set incorrectly.

[ROCm] Use specific amdgpu version for EL8 systems

We were always installing the latest driver versions
but this had some side effects when yum would try
to download index files from a URL with changing content.

[ROCm] Fix formatting on python files

Reformatted with black
 2024-09-30 12:39:51 -05:00
Mathew Odden
a9e54b3e0a Add docker builds for ubu22 and 24 2024-08-27 16:40:37 -05:00
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