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

280 Commits

Author SHA1 Message Date
Gary Miguel
dc908b4843 Update installation instructions 
Apple GPUs and Mac x86_64 is a non-existent combination.
Mac x86_64 with AMD GPU is supported.

It's a bit of a confusing situation so hard to summarize, but hopefully this is more accurate and less confusing

Fixes: #24408
 2024-10-21 10:20:09 -07:00
Ruturaj4
3c3b08dfd6 [ROCm] Fix README.md to update AMD JAX installation instructions 2024-10-16 17:15:32 -05:00
Jake VanderPlas
a8744e38d3 README: fix badge href links 2024-10-02 10:09:00 -07:00
Michael Hudgins
d4d1518c3d Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax 
PiperOrigin-RevId: 676843138
 2024-09-20 07:52:33 -07:00
Matthew Johnson
95cb347d04 trying on another readme header 2024-09-10 17:15:29 +00:00
Matthew Johnson
8ab503158b tweak readme title to be more about what jax can do for you, dear user 
we should rewrite this whole readme...
 2024-09-07 18:09:19 +00:00
Roy Frostig
c140103a02 new README tagline 2024-08-29 13:23:14 -07:00
Roy Frostig
371935cc10 update README and several docs to typed RNG keys 2024-08-11 08:09:47 -07:00
Peter Hawkins
3735a485b0 Remove [cpu] extra from installation instructions. 
`pip install jax` should do the right thing by default now.
 2024-06-18 11:12:23 -04:00
Peter Hawkins
c71298c2ee Temporarily reinstate the jax[cpu] in the install instructions. 
This is still needed until the next release.
 2024-06-13 17:00:20 -04:00
Peter Hawkins
b13733c13f Update JAX dependencies, extras, and documentation for plugins. 
* Make jaxlib a direct dependency of jax.
* Remove mentions of monolithic CUDA installations from the JAX documentation.
* Drop the cuda12_pip extra and the cudnn version specific extras.
* Add a with_cuda extra to the jax-cuda12-plugin package, use it in jax's setup.py. This allows us to specify cuda extras in one place.
* Make a few small doc improvements.
 2024-06-13 11:36:23 -04:00
Cristian Garcia
ded86e6e6b
Add reference to Flax NNX 
Adds a reference to Flax NNX in the README.
 2024-05-27 10:27:57 +01:00
Jieying Luo
5ba56bb075 Recommend the plugin in the CUDA installation instructions. 
PiperOrigin-RevId: 631555876
 2024-05-07 14:47:39 -07:00
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