mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-22 06:46:04 +00:00
Code Issues Packages Projects Releases Wiki Activity
18,624 Commits 508 Branches 263 Tags
Commit Graph

290 Commits

Author SHA1 Message Date
Antonio Sanchez
873cffc776 Use TSL's import of DUCC. 
This is a necessary first step before adding DUCC support to XLA,
otherwise the JAX tests in the XLA repo pull from JAX's copy,
which has slightly different build rules.

PiperOrigin-RevId: 576880208
 2023-10-26 08:26:56 -07:00
Peter Hawkins
70b7d50181 Switch jaxlib to use nanobind instead of pybind11. 
nanobind has a number of advantages (https://nanobind.readthedocs.io/en/latest/why.html), notably speed of compilation and dispatch, but the main reason to do this for these bindings is because nanobind can target the Python Stable ABI starting with Python 3.12. This means that we will not need to ship per-Python version CUDA plugins starting with Python 3.12.

PiperOrigin-RevId: 559898790
 2023-08-24 16:07:56 -07:00
Peter Hawkins
d6e06f4476 Move the XLA commit out of the top-level JAX WORKSPACE file and into a separate .bzl file. 
No functional changes intended.

PiperOrigin-RevId: 556906943
 2023-08-14 14:09:45 -07:00
Skye Wanderman-Milne
e132a0e5d5 Slightly downgrade xla version to avoid PJRT C API incompat 2023-07-27 14:05:14 -07:00
Skye Wanderman-Milne
0b24b2ba6a Update WORKSPACE and setup.py in preparation for jax/jaxlib 0.4.14 release 2023-07-27 13:35:04 -07:00
Peter Hawkins
0cd025348d Update XLA version to fix build failure. 2023-07-27 12:38:51 +00:00
Sharad Vikram
3baa6e7a89 Enable building jaxlib w/ Mosaic 
PiperOrigin-RevId: 551159246
 2023-07-26 03:59:30 -07:00
Peter Hawkins
487b640acf Jax 0.4.13 release. 2023-06-22 14:59:36 -07:00
Skye Wanderman-Milne
8223e452f2 Update WORKSPACE and setup.py for jax + jaxlib 0.4.12 release 2023-06-08 13:58:09 -07:00
Yash Katariya
48ad9a6f3e Start jax and jaxlib 0.4.11 release 
PiperOrigin-RevId: 536860076
 2023-05-31 16:48:52 -07:00
Sharad Vikram
1279418ce5 Link in CUDA runtime for triton in jaxlib 
PiperOrigin-RevId: 535708416
 2023-05-26 14:02:16 -07:00
Skye Wanderman-Milne
82bbeef519 Update setup.py, WORKSPACE, and CHANGELOG for jax/jaxlib 0.4.10 release 2023-05-11 14:46:06 -07:00
Skye Wanderman-Milne
5bcd9dcc46 Update WORKSPACE and setup.py in preparation for jax/jaxlib 0.4.9 release, take 2 2023-05-09 14:49:54 -07:00
Skye Wanderman-Milne
f28b20175f Update WORKSPACE and setup.py in preparation for jax/jaxlib 0.4.9 release 2023-05-04 14:38:46 -07:00
Peter Hawkins
51c40f04b9 Bump XLA version. 2023-05-03 18:52:04 +00:00
Yash Katariya
ae4f1fcb66 Update the commit in workspace too 
PiperOrigin-RevId: 519797427
 2023-03-27 12:30:18 -07:00
Yash Katariya
e9cac5eb47 Prepare for jax and jaxlib 0.4.7 release 
PiperOrigin-RevId: 519785176
 2023-03-27 11:45:22 -07:00
Peter Hawkins
172a831219 Switch JAX to use the OpenXLA repository. 2023-03-13 18:38:26 +00:00
Skye Wanderman-Milne
1aa08fd4a0 Update WORKSPACE and setup.py for jax/jaxlib 0.4.6 release 2023-03-09 10:44:58 -08:00
Yash Katariya
58e46b48e6 Prepare for jax and jaxlib 0.4.4 release 
PiperOrigin-RevId: 510152471
 2023-02-16 08:37:15 -08:00
Skye Wanderman-Milne
8ab158574d Update WORKSPACE and setup.py for jax/jaxlib 0.4.3 release 2023-02-07 15:45:28 -08:00
Skye Wanderman-Milne
9a7f29a9fb Update WORKSPACE for jaxlib 0.4.2 release (again) 
This fixes a segfault that was affecting the previous 0.4.2 candidate.
 2023-01-24 18:07:11 +00:00
Skye Wanderman-Milne
3f4bd5f449 Updates for jax + jaxlib 0.4.2 release 2023-01-20 19:04:46 +00:00
Yash Katariya
c4d590b1b6 Update values for release 0.4.1 
PiperOrigin-RevId: 494889744
 2022-12-12 19:04:38 -08:00
Yash Katariya
0118f8d568 Prepare for jax and jaxlib 0.4.0 release 
PiperOrigin-RevId: 493733609
 2022-12-07 16:02:24 -08:00
Jake VanderPlas
e66fe1dff4 Update XLA commit. 
Fixes build error: no such target '@org_tensorflow//tensorflow/compiler/xla/mlir_hlo:python/MlirHloModule.cc'
 2022-11-16 15:51:28 -08:00
Yash Katariya
50f7512a6e
Update WORKSPACE for jaxlib release 2022-11-14 16:37:09 -08:00
Yash Katariya
2e384ce58f Prepare for release of jax and jaxlib 0.3.24 
PiperOrigin-RevId: 485985460
 2022-11-03 15:13:23 -07:00
Skye Wanderman-Milne
63be2201aa Update WORKSPACE and setup.py for jaxlib 0.3.22 release 2022-10-11 17:55:05 +00:00
Peter Hawkins
8d8643664c jax/jaxlib 0.3.20 release candidate. 2022-09-28 13:33:52 +00:00
Peter Hawkins
bcd36d8eb2 Jax and jaxlib 0.3.18 release candidate. 2022-09-26 14:10:57 +00:00
Peter Hawkins
bec9d2e168 Bump XLA version. 2022-09-19 13:17:46 +00:00
jax authors
45764ea9a8 Merge pull request #12125 from hawkinsp:xla 
PiperOrigin-RevId: 470295638
 2022-08-26 11:38:03 -07:00
Peter Hawkins
12cb55d2f0 Bump XLA version. 2022-08-26 18:27:59 +00:00
Peter Hawkins
b63801b4db Fixes for PocketFFT->ducc migration. 
* Rename modules from pocketfft to ducc.
* Fix up strides at their generation point rather than where they are
  consumed.
 2022-08-26 14:30:03 +00:00
Peter Hawkins
6128cb7aa0 Update XLA commit. 2022-08-18 20:17:37 +00:00
Peter Hawkins
87db8fc5f6 Update XLA commit. 
Fixes build error:
Label
 '@org_tensorflow//tensorflow/tsl/platform/default:build_config.bzl' is
 invalid because 'tensorflow/tsl/platform/default' is not a package.
 2022-08-17 15:01:43 +00:00
Adam Paszke
6ed2d22566 Update WORKSPACE 
to include the addition of use_global_device_ids in AllReduceOp.
 2022-08-17 07:30:54 +00:00
Peter Hawkins
5603cb31d4 Bump XLA version. 
Fixes https://github.com/google/jax/issues/9771
 2022-08-15 21:12:52 +00:00
Yash Katariya
e982dd105f
Update the TF commit hash 2022-08-05 10:51:21 -07:00
Peter Hawkins
96cb47df32 Update XLA. 2022-08-01 12:10:27 +00:00
Skye Wanderman-Milne
26fbeb6e2a Update WORKSPACE and libtpu version for jaxlib 0.3.15, take 3 2022-07-22 11:41:39 -07:00
Skye Wanderman-Milne
568cedba8d
Update WORKSPACE for 0.3.15 release, take 2 2022-07-20 17:23:52 -07:00
Skye Wanderman-Milne
9149c38e1e Update WORKSPACE and setup.py in preparation for 0.3.15 jax/jaxlib release 2022-07-14 10:12:59 -07:00
Peter Hawkins
7c49864fdf Symlink xla_client and xla_extension into jaxlib rather than copying them into place in the wheel build. 
Change in preparation for allowing JAX tests to run under Bazel.

Remove code to patch paths in xla_client.py in the wheel build script; the patch is no longer used.

PiperOrigin-RevId: 458522398
 2022-07-01 12:31:42 -07:00
Peter Hawkins
f4ddd3ef88 Update XLA. 2022-06-27 09:14:28 -04:00
Sharad Vikram
217d898124 Update TF version for jaxlib build 2022-06-21 14:34:22 -07:00
Sharad Vikram
9bd1bd67e0 Update versions for jax/jaxlib release 2022-06-21 12:57:28 -07:00
Peter Hawkins
5b7c485ade Update XLA pin. 2022-06-15 09:24:57 -04:00
Peter Hawkins
1e504b36e2 Improve the comment in WORKSPACE describing how to override the TF repository. 2022-06-14 22:25:39 -04:00