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

44 Commits

Author SHA1 Message Date
Peter Hawkins
85add667ed Change documentation to recommend libtpu from pypi instead of GCS. 2025-02-24 17:50:29 -05:00
Jake VanderPlas
7eacce0b83 DOC: update osx x86 entry in installation grid 2025-01-27 11:18:33 -08:00
Robert Dyro
aa9cea0a55 Add quotes to pip commands in docs around option install for zsh 
PiperOrigin-RevId: 716358130
 2025-01-16 13:30:43 -08:00
Ruturaj4
fefebebea4 Update documentation and add setup.py pypi bindings 2024-12-19 10:21:28 -06:00
Silvio Traversaro
09309e6452
Update conda-forge installation docs after CUDA 12 upgrade 2024-12-10 12:11:28 +01: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
Li-Jesse-Jiaze
5e1366c4ce Fix #24661: Add zsh support to conda install documentation 2024-11-01 17:57:18 +01:00
minigoel
68428488c8 Add a link to Intel plugin for JAX 2024-10-28 10:47:59 -07:00
Peter Hawkins
e4f3f8f064 Use libtpu releases rather than libtpu-nightly for jax[tpu]. 
PiperOrigin-RevId: 688632409
 2024-10-22 11:47:07 -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
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
Yash Katariya
8ab66c8103 Fix the TPU and GPU nightly install instructions. 
PiperOrigin-RevId: 675233702
 2024-09-16 11:46:58 -07:00
Peter Hawkins
b975592478 Change nightly install commands to include all packages. 
pip doesn't update transitive dependencies, and we probably want the latest versions of everything when installing a nightly.
 2024-09-09 14:38:41 -04:00
Roy Frostig
672a013b3a docs: prefer "summary" to "tl;dr" 
This was common and we typically just mean "summary."
 2024-08-28 11:04:14 -07:00
Roy Frostig
3c223cd253 docs: tidy up titles and headings 
This shortens some titles and makes them more consistent. It also
removes "JAX" from several titles ("in JAX", "for JAX", "JAX's",
etc.). Since these are JAX docs, that ought to be clear from context.
 2024-08-13 11:53:57 -07:00
Peter Hawkins
d1c0d993fc Bump the minimum CUDNN version to v9.1. 
This actually was already the minimum version since we build with that version, but we needed to tighten the constraints.

Also in passing, drop mentions of CUDA builds from the Windows build instructions. jaxlib hasn't built with CUDA enabled on Windows for a very long time, so it's probably best we just don't mention it.

PiperOrigin-RevId: 657225917
 2024-07-29 09:28:47 -07:00
jax authors
00528b9858 libdevice.10.bc is removed from JAX wheels bundle. 
The recommended source of JAX wheels is `pip`, and NVIDIA dependencies are installed automatically when JAX is installed via `pip install`. `libdevice` gets installed from `nvidia-cuda-nvcc-cu12` package.

PiperOrigin-RevId: 647328834
 2024-06-27 08:35:59 -07:00
jax authors
fc1e1d4a65 Add freshness metablock to JAX OSS docs. 
PiperOrigin-RevId: 645508135
 2024-06-21 14:50:49 -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
jax authors
2f22d3abfd Merge pull request #21901 from jjyyxx:patch-1 
PiperOrigin-RevId: 644061544
 2024-06-17 10:44:50 -07:00
Yuxuan Jiang
cd23b2c82c
Fix CUDNN requirement inconsistency in installation.md 2024-06-16 14:38:05 +08:00
Erick Matsen
b585eceeb0
Fixing installation docs: needed quotes; markdown fix 2024-06-15 05:12:21 -07:00
Peter Hawkins
d5844b7cd3 Update nightly installation instructions to reflect new jax/jaxlib dependency structure. 2024-06-14 15:48:31 -04:00
Yash Katariya
8f315c212a Merge all nightly wheels into jax_nightly_releases.html since we have dropped monolithic jaxlib CUDA wheels 
PiperOrigin-RevId: 643331639
 2024-06-14 07:01:44 -07: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
Jieying Luo
5ba56bb075 Recommend the plugin in the CUDA installation instructions. 
PiperOrigin-RevId: 631555876
 2024-05-07 14:47:39 -07:00
Jake VanderPlas
a4ccb354f8 DOC: update installation guide 2024-04-16 12:46:36 -07:00
Peter Hawkins
478cfa9944 Add an upper bound on JAX's CUDNN version constraint. 
Major releases of CUDNN break ABI compatibility, so we cannot allow new major versions.

PiperOrigin-RevId: 620030416
 2024-03-28 13:00:36 -07:00
jax authors
0be07e6aec Remove support for CUDA 11. 
Pin minimal required versions for CUDA to 12.1.

Reverts 910a31d7b7510e3375718ab1ea0d38df7bd2c0d5

PiperOrigin-RevId: 618911489
 2024-03-25 11:46:39 -07:00
jax authors
910a31d7b7 Reverts bed4f65438a62777ed100ecec2b0eb3f7cf87a0e 
PiperOrigin-RevId: 618249855
 2024-03-22 12:10:53 -07:00
jax authors
bed4f65438 Remove support for CUDA 11. 
Pin minimal required versions for CUDA to 12.1.

PiperOrigin-RevId: 618195554
 2024-03-22 09:05:39 -07:00
Yash Katariya
4bb0a5a204 Add --pre to the libtpu nightly install instructions. 
PiperOrigin-RevId: 609702726
 2024-02-23 05:52:59 -08:00
Yash Katariya
73e3dedf9d Update changelog and install doc to mention cuda 12.3 switch 
PiperOrigin-RevId: 605473026
 2024-02-08 17:21:56 -08:00
Jake VanderPlas
c785eeb4e5 DOC: add a note about installing older GPU wheels 2024-01-12 11:12:39 -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
Jake VanderPlas
3567ddb6df Document installation of older jaxlib wheels 2023-11-03 12:34:22 -07:00
Yash Katariya
87c0561093 Add nightly install instructions to docs 
PiperOrigin-RevId: 576990265
 2023-10-26 14:36:39 -07:00
Peter Hawkins
04bee99283 Update references to aarch64 support in the README. 2023-10-10 12:36:29 -04:00
Peter Hawkins
8c4d020db9 Improve CUDA install documentation. 
Mention NCCL as a dependency, since it will be required by the next jaxlib release.
Mention LD_LIBRARY_PATH and PATH as how one overrides the CUDA installation for local installs.

Fixes #17831
 2023-10-06 14:36:29 -04:00
Peter Hawkins
b7dfde8d87 Add notes about the new CUDA version restrictions to the changelog and installation instructions. 2023-09-27 15:56:47 -04: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