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

267 Commits

Author SHA1 Message Date
Jake VanderPlas
10ed827fe9 DOC: replace old tutorials with new content 2024-05-03 12:20:06 -07:00
Jake VanderPlas
650fd780f7 README: improve TPU accuracy discussion 2024-04-02 09:16:57 -07:00
Jake VanderPlas
72b2321750 DOC: update JAX quick description 2024-02-26 14:14:11 -08:00
Skye Wanderman-Milne
5d26c307ce Revert "Recommend the plugin in the CUDA installation instructions." 
This reverts commit 78bc4f5ced41f7f96a70da769e7c3170dc2b3161.

GPU dlpack is broken with the new plugin. Recommend working GPU jaxlib until we can fix.
 2023-12-14 17:23:18 -08:00
Peter Hawkins
78bc4f5ced Recommend the plugin in the CUDA installation instructions. 2023-12-14 17:21:30 -05:00
Peter Hawkins
04bee99283 Update references to aarch64 support in the README. 2023-10-10 12:36:29 -04:00
bb515
b22a29fa16
Update README.md so that CPU install instruction works 
Fix readme
current `pip install -U jax[cpu]` has no matches in pip `no matches found: jax[cpu]`.
Corrected install instruction to `pip install -U "jax[cpu]"` which successfully installs cpu version of JAX via pip.
 2023-10-10 12:03:16 +01:00
Matthew Johnson
dddbe43444
Update README.md 
Co-authored-by: 8bitmp3 <19637339+8bitmp3@users.noreply.github.com>
 2023-10-03 12:33:41 -07:00
Matthew Johnson
3b0c7892f1 tweak readme to call out broader numerical computing 
at the same time, it makes sense to highlight large-scale machine learning as a focus
 2023-09-29 08:59:27 -07:00
Peter Hawkins
fb97f48870 Update install instructions for AMD (ROCm) GPUs to point to AMD's docker container. 2023-09-26 09:52:45 -04:00
Peter Hawkins
8cbd58a9f0 Improve installation instructions. 
Simplify the common case installation instructions on the JAX readme.

Move less-commonly used or more complicated options to the installation page of the JAX docs.
 2023-09-22 17:13:47 -04:00
cabbagepatchman
125343b7f9
Update README.md 2023-09-12 20:39:06 -07:00
Patrick Kidger
559a27e6c3
Update README.md 
Removed Haiku (no longer supported by DeepMind); added mention of Equinox (supported by X).
 2023-09-07 09:55:56 +01:00
jax authors
d03187e305 Merge pull request #17060 from nouiz:rosetta_doc 
PiperOrigin-RevId: 558886922
 2023-08-21 13:37:24 -07:00
8bitmp3
d300c5c692
Update JAX README.md - remove jax.example_libraries, stax 2023-08-17 23:42:19 +00:00
Frederic Bastien
d0059b70be typo 2023-08-17 07:46:35 -07:00
Frederic Bastien
cff08d1057 Small doc update following code review 2023-08-16 09:23:58 -07:00
Frederic Bastien
f212634454 Add a link to JAX Toolbox containers 2023-08-10 12:38:25 -07:00
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