mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 20:06:05 +00:00
Code Issues Packages Projects Releases Wiki Activity
26,442 Commits 502 Branches 263 Tags
Commit Graph

291 Commits

Author SHA1 Message Date
Charles Hofer
7a6940bb7b Trivial change for CI 2025-03-14 22:45:17 +00:00
Charles Hofer
db8ba1b598 Change to run CI 2025-03-12 17:06:35 +00:00
Peter Hawkins
85add667ed Change documentation to recommend libtpu from pypi instead of GCS. 2025-02-24 17:50:29 -05:00
Ruturaj4
fefebebea4 Update documentation and add setup.py pypi bindings 2024-12-19 10:21:28 -06:00
George Necula
9f6cf62a48 Minor change in the README, remove "expect bugs" 2024-12-11 14:50:14 +01:00
Emily Fertig
225a2a5f8b Consolidate material on PRNGs and add a short summary to Key Concepts. 2024-11-15 14:44:57 -08:00
jax authors
4511f0c66b Merge pull request #24862 from emilyfertig:emilyaf-control-flow-tutorial 
PiperOrigin-RevId: 696692588
 2024-11-14 16:50:14 -08:00
jax authors
8e292122b7 Merge pull request #24567 from Intel-tensorflow:minigoel/intel-plugin 
PiperOrigin-RevId: 696677564
 2024-11-14 15:52:38 -08:00
Emily Fertig
e6f6a8af8d Move Control Flow text from Sharp Bits into its own tutorial. 2024-11-14 11:07:52 -08:00
8bitmp3
b976b1ab27 Update details on JAX libraries in JAX README.md 2024-11-05 00:37:12 +00:00
minigoel
68428488c8 Add a link to Intel plugin for JAX 2024-10-28 10:47:59 -07:00
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