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
9,481 Commits 502 Branches 263 Tags
Commit Graph

213 Commits

Author SHA1 Message Date
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
Qiao Zhang
a22841b6bb Bump jaxlib ver to 0.1.68. 2021-06-23 12:37:56 -07:00
George Necula
6a48c60a72 Rename master to main in embedded links. 
Tried to avoid the change on external links to repos that
have not yet renamed master.
 2021-06-18 10:00:01 +03:00
Skye Wanderman-Milne
63dbb99a66 Update README, etc. for jaxlib 0.1.67 release 2021-05-17 17:48:46 -07:00
Peter Hawkins
f88e45295f Update installation instructions for CUDA wheels. 2021-05-12 09:24:31 -04:00
Qiao Zhang
528d5bbb11 Update README etc for jaxlib 0.1.66 release. 2021-05-11 16:49:32 -07:00
Skye Wanderman-Milne
f8f373466c Update README, etc. for jaxlib 0.1.65 release 2021-04-07 17:51:20 -07:00
Lena Martens
25a2e1a66a
Fix the changelog link in the README. 2021-03-25 19:15:45 +00:00
Skye Wanderman-Milne
0cbe2c1c05 Update README, etc. for jaxlib 0.1.64 release 2021-03-18 16:11:40 -07:00
Skye Wanderman-Milne
757247b791 Update README, etc. for jaxlib 0.1.63 release 2021-03-17 10:14:52 -07:00
Tom Hennigan
4f74b3391c
Update README.md 
Co-authored-by: 8bitmp3 <19637339+8bitmp3@users.noreply.github.com>
 2021-03-12 16:38:29 +00:00
Tom Hennigan
0fd83a30be Link to DeepMind JAX blog post. 2021-03-12 12:01:20 +00:00
Skye Wanderman-Milne
f06bb9a7f4 Update jaxlib version etc. 2021-03-09 17:55:40 -08:00
Jake VanderPlas
749ad95514 DOC: add transformations doc to HTML & reorganize contents 2021-03-08 16:25:04 -08:00
Jake VanderPlas
24e920e3f5 change colab.sandbox->colab.research 2021-02-25 10:29:43 -08:00
Skye Wanderman-Milne
7a67b974ac jaxlib version bump etc. 2021-02-12 09:42:04 -08:00
Peter Hawkins
13f3819054 Update README.md for jaxlib 0.1.60. 
Bump jaxlib version to 0.1.61 and update changelog.

Change jaxlib numpy version limit to >=1.16 for next release. Releases older than 1.16 are deprecated per NEP 00029. Reenable NumPy 1.20.

Bump minimum jaxlib version to 0.1.60.
 2021-02-03 20:44:01 -05:00
Peter Hawkins
4fdd029b3e Clarify the CuDNN versions expected by jaxlib wheels. 2021-01-20 10:00:19 -05:00
Skye Wanderman-Milne
7c2454e969 Update jaxlib version, minimum jaxlib version, readme, and changelog. 
Bumping the min jaxlib version to support https://github.com/google/jax/pull/5213.
 2021-01-15 12:56:08 -08:00
Peter Hawkins
5747c91564 Update CUDA installation instructions. 2021-01-13 10:48:43 -05:00
jax authors
783e1b7c89 Merge pull request #5236 from bileschi:patch-1 
PiperOrigin-RevId: 350975202
 2021-01-09 20:43:00 -08:00
Stanley Bileschi
c1904ea7a7 fixing. Renaming per reviewer feedback. 2021-01-06 11:26:46 -05:00
John Aslanides
da4e27315d Add PyPI version badge to README. 2020-12-30 22:48:24 +00:00
Stanley Bileschi
80173aad62
(docfix) Typo in README.md 
`predict` function should apply nonlinearity to output
 2020-12-21 00:35:22 -05:00
Peter Hawkins
c83343d80b Add a few notes about Windows builds and a link to the build instructions. 2020-11-30 20:13:01 -05:00
jax authors
6cdd956a74 Merge pull request #5042 from redwrasse:redwrasse-patch-1 
PiperOrigin-RevId: 344799393
 2020-11-30 06:17:58 -08:00
redwrasse
06957e8caa
fix readme 
fix readme
 2020-11-29 18:51:24 -07:00
Matthew Johnson
f180ea86d4
fix capitalization 2020-11-16 19:03:15 -08:00
Matthew Johnson
c2a98c1820
update jax repository citation 
Add all full-time JAX team members.
 2020-11-16 19:01:00 -08:00
Skye Wanderman-Milne
36799f5007 README etc. updates for new jaxlib release 2020-11-12 13:48:44 -08:00
mariosasko
68ef05be46 Fix broken link in readme 2020-11-09 00:33:04 +01:00
Matthew Johnson
a7d1963bc8 update neural network libraries text 
Co-authored-by: George Necula <necula@google.com>
 2020-10-30 08:42:04 -07:00
Matthew Johnson
a4ce2813c8 link to nn libraries in readme 
Co-authored-by: Skye Wanderman-Milne <skyewm@google.com>
Co-authored-by: James Bradbury <jekbradbury@google.com>
 2020-10-30 08:39:59 -07:00
Skye Wanderman-Milne
a082169642 Update README and jaxlib.__version__ for new jaxlib release 2020-10-20 11:46:48 -07:00
Alexey Radul
537427aae2 Copybara import of the project: 
--
138105a9ea44e7a8c3ce575a4e51b7ed51518d41 by Skye Wanderman-Milne <skyewm@google.com>:

Update README, CHANGELOG, and jaxlib.__version__ for new jaxlib release

PiperOrigin-RevId: 338063494
 2020-10-20 08:25:23 -07:00
Skye Wanderman-Milne
138105a9ea Update README, CHANGELOG, and jaxlib.__version__ for new jaxlib release 2020-10-14 13:32:56 -07:00
Skye Wanderman-Milne
cacb01753a Use local version identifiers to distribute cuda jaxlib wheels. 
This change:

* Updates our jaxlib build scripts to add `+cudaXXX` to the wheel
  version, where XXX is the CUDA version number (e.g. `110`). nocuda
  builds remain unchanged and do not have this extra identifier.

* Adds `generate_release_index.py`, which writes an html page that pip
  can use to find the cuda wheels. (I based this format off of
  wheel PyTorch's index).

* Updates the README to use the new local version identifier + wheel
  index.

The end result is that the command to install cuda wheels is now much
simpler.

I manually made copies of the latest jaxlib 0.1.55 wheels that have
the local version identifiers, so the new installation commands
already work (as well as the old ones, until the next jaxlib release
using the new tooling).

Fow now, I put the html index to the GCP bucket with the wheels. We
can move it to a prettier URL if/when we have one.
 2020-10-09 13:47:54 -07:00