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

249 Commits

Author SHA1 Message Date
Peter Hawkins
a1a5159fbf Note that the MSVC studio 2019 redistributable is required for JAX on Windows. 
Issue #16664
 2023-07-11 12:45:13 -04:00
Peter Hawkins
a1de687382 Update README to mention Windows support. 2023-06-23 16:17:18 -04:00
Peter Hawkins
487b640acf Jax 0.4.13 release. 2023-06-22 14:59:36 -07:00
Peter Hawkins
85a84fdee8 Add a link to the Apple Metal plugin to the JAX README. 
Remove references to cloud TPU on colab, since support was dropped in JAX 0.4. Users should use TPUs via Kaggle or via Cloud TPU VMs.
 2023-06-22 09:09:24 -04:00
Yash Katariya
be68ca65a5 Update the cuda 12 wheel name to cudnn 8.9 since that is what it is built with. Fixes https://github.com/google/jax/issues/16275 
PiperOrigin-RevId: 538337302
 2023-06-06 17:38:20 -07:00
Peter Hawkins
69cf67f252 Bump the minimum CUDNN version for CUDA 12 wheels to 8.9. 2023-05-26 10:04:34 -04:00
Michael Deistler
5f1952df4d
fix typo 2023-05-24 10:43:03 +02:00
Peter Hawkins
ee4bc2aa17 Add a note that Kepler GPUs are not supported to the README. 2023-05-19 11:50:45 -04:00
Peter Hawkins
46d30ba29e Drop mentions of the CUDA 11.4/CUDNN 8.2 wheel. 
This wheel is no longer shipped as part of jaxlib releases; please upgrade if you were using this.
 2023-05-12 10:41:48 -04:00
Peter Hawkins
775f404f31 Update the CUDA installation instructions. 2023-03-28 21:46:07 -04:00
jax authors
d19e60ea07 Merge pull request #15228 from canyon289:patch-2 
PiperOrigin-RevId: 519742908
 2023-03-27 09:30:07 -07:00
Ravin Kumar
08a8a5e812
Fix hessian llnk 2023-03-26 19:21:52 -07:00
Peter Hawkins
b62f114524 Add support for using pip-installed CUDA wheels. 
Add a currently undocumented jax[cuda11_pip] and jax[cuda12_pip] that depend on the pip CUDA wheels.
Add a currently undocumented jax[cuda11_local] and jax[cuda12_local] that avoid the CUDA wheel dependency.
 2023-03-26 12:35:00 +00:00
Jake VanderPlas
e3444a8d42 README: improve Colab TPU installation discussion 2023-03-15 08:54:23 -07:00
Skye Wanderman-Milne
56b237cfbc Update Cloud TPU install command to be simpler. 
We used to need the extra stuff for a very old Cloud TPU VM image, but we don't anymore.
 2023-02-27 23:15:19 +00:00
Peter Hawkins
eb1710f84a Update conda GPU install instructions. 2023-02-23 08:37:10 -05:00
Owen Lockwood
2230225956 Update README.md 
Update README.md
 2023-01-06 15:01:43 -07:00
Peter Hawkins
7bd6cc13e0 Note CUDA 11.8 is required for Hopper and Ada Lovelace series GPUs. 
Fixes https://github.com/google/jax/issues/13570
 2022-12-12 11:03:48 -05: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
79c8d6679f README: fix badge URL 2022-12-07 08:48:40 -08:00
Jimmy
e3cd9671b9 Add x86_64 dependency note to pip installation 
Currently non x86_64 linux architectures are not supported, see #7097 for request to change this. This can lead to installation confusion, as jax will install, but jaxlib will not. For example see #12307. This adds a note to the install sections for the relevant pip wheels.
 2022-11-04 12:11:11 -07:00
Peter Hawkins
8f1a346198 Recommend using the cuda-nvcc package from the "nvidia" conda channel to fetch ptxas. 2022-08-16 13:34:07 -04:00
Peter Hawkins
026e760767 Point to the conda-forge jaxlib wheels in the JAX readme. 2022-08-15 20:34:09 +00:00
Peter Hawkins
a2c21958a5 Document multiprocess GPU support. 
Fixes #2731
 2022-08-09 11:31:05 -04:00
Matthew Johnson
3ab8637a73 Update README discussion of Windows support: 
* right next to the pip installation instructions, mention they don't work for Windows;
* add a link to #5795 for an unofficial discussion of Windows native support
 2022-07-19 17:12:40 -07:00
Yash Katariya
e780a402e8 Split cuda and nocuda index files into their own 
PiperOrigin-RevId: 455463113
 2022-06-16 14:07:55 -07:00
Roy Frostig
531ab3fb18 add CITATION.bib and bump version to latest release 
PiperOrigin-RevId: 453519227
 2022-06-07 14:00:21 -07:00
Filippo Vicentini
a076d349d9
Make CUDA install instructions copy-pasteablee 2022-05-11 17:57:26 +02:00
Dean Biskup
554b02e978
Add quotes to install commands in README 2022-04-19 12:23:52 -05:00
Peter Hawkins
38ea5a6bc0 Copybara import of the project: 
--
391dea76bc8fe264cf26ec93d42147f87847894d by Peter Hawkins <phawkins@google.com>:

Update version numbers after jax/jaxlib 0.3.7 release.

COPYBARA_INTEGRATE_REVIEW=https://github.com/google/jax/pull/10324 from hawkinsp:jaxlib 391dea76bc8fe264cf26ec93d42147f87847894d
PiperOrigin-RevId: 442311051
 2022-04-16 22:37:09 -07:00
Yazeed Alaudah
4a8de96e88
Update README.md to fix link to prng.md 2022-02-20 09:46:49 -06:00
Peter Hawkins
d704c151fa Clarify the NVidia driver version requirements. 2022-02-17 14:37:29 -05:00
Peter Hawkins
384042d0e7 Clarify the required NVidia driver versions in the JAX readme instructions. 2022-01-28 11:35:34 -05:00
Jake VanderPlas
0a44aa69df README: clarify installation instructions for Colab TPU 2022-01-26 11:10:58 -08:00
Matthew Johnson
57f4fc9bbb update readme to focus on most active libraries 2021-11-15 13:47:13 -08:00
Brian Hulette
6cfbb890e6 Fix broken link to experimental/README.md 2021-10-20 13:01:26 -07:00
Roy Frostig
623c201054 [JAX] move example libraries from jax.experimental into jax.example_libraries 
The `jax.experimental.stax` and `jax.experimental.optimizers` modules are standalone examples libraries. By contrast, the remaining modules in `jax.experimental` are experimental features of the JAX core system. This change moves the two example libraries, and the README that describes them, to `jax.example_libraries` to reflect this distinction.

PiperOrigin-RevId: 404405186
 2021-10-19 17:30:45 -07:00
Yash Katariya
ee752b32f7 Use cuda11_cudnn82 instead of cuda=11,cudnn=82 because the latter one is a syntax error 
PiperOrigin-RevId: 404240654
 2021-10-19 06:24:53 -07:00
Yash Katariya
4d8bce1b85 Add a default cuda installation path and more explicit installation paths for CUDA jaxlib. 
```
# Installs Cuda 11 with Cudnn 8.2
$ pip install jax[cuda] -f https://storage.googleapis.com/jax-releases/jax_releases.html

$ pip install jax[cuda=11,cudnn=82] -f https://storage.googleapis.com/jax-releases/jax_releases.html

$ pip install jax[cuda=11,cudnn=805] -f https://storage.googleapis.com/jax-releases/jax_releases.html
```

PiperOrigin-RevId: 404134291
 2021-10-18 19:56:22 -07:00
Yash Katariya
66a4a9ff3f Remove 10.2 cuda support 
PiperOrigin-RevId: 402707900
 2021-10-12 18:44:07 -07:00
Jake VanderPlas
d0e21072ec DOC: add brief comments to example in README 2021-09-14 11:47:19 -07:00
josephrocca
507cc99c84 Add a note about Google Colab setup 
User must run `jax.tools.colab_tpu.setup_tpu()`
 2021-09-13 21:34:02 -07:00
Yash Katariya
b603b34d5f Update the readme to reflect that Cuda 10.1 and cuda 11.0 support is dropped. 
PiperOrigin-RevId: 394772160
 2021-09-03 15:56:47 -07:00
Jake VanderPlas
e46200e123 README: mention that CUDA installation is for linux only 2021-08-28 08:15:21 -07:00
yashkatariya
1f13565b70 Remove opensource CS link since JAX is not indexed 2021-08-12 11:12:35 -07:00
elliotwaite
7392a57b75 DOC: many small fixes 2021-08-04 16:55:13 -07:00
Shawn Presser
f800029569
Fix broken links in README.md 2021-07-25 09:40:24 -05:00
Jamie Townsend
6ca775b10a
Add TPU precision details to README gotchas 2021-07-16 12:24:19 +02:00
Jake VanderPlas
5e2681939a DOC: update installation instructions for compatibility with zsh 2021-06-28 12:57:24 -07:00
Jake VanderPlas
3c727ab604 DOC: update install instructions for GPU & TPU 2021-06-24 11:22:12 -07:00