mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-15 19:36:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Update details on JAX libraries in JAX README.md

Browse Source
This commit is contained in:
8bitmp3 2024-11-04 22:35:26 +00:00
parent 38b4d00100
commit b976b1ab27
1 changed files with 10 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
README.md
View File

@ -411,23 +411,18 @@ community-supported conda build, and answers to some frequently-asked questions.
## Neural network libraries ## Neural network libraries
Multiple Google research groups develop and share libraries for training neural Multiple Google research groups at Google DeepMind and Alphabet develop and share libraries
networks in JAX. If you want a fully featured library for neural network for training neural networks in JAX. If you want a fully featured library for neural network
training with examples and how-to guides, try training with examples and how-to guides, try
[Flax](https://github.com/google/flax). Check out the new [NNX](https://flax.readthedocs.io/en/latest/nnx/index.html) API for a [Flax](https://github.com/google/flax) and its [documentation site](https://flax.readthedocs.io/en/latest/nnx/index.html).
simplified development experience.
Google X maintains the neural network library Check out the [JAX Ecosystem section](https://jax.readthedocs.io/en/latest/#ecosystem)
[Equinox](https://github.com/patrick-kidger/equinox). This is used as the on the JAX documentation site for a list of JAX-based network libraries, which includes
foundation for several other libraries in the JAX ecosystem. [Optax](https://github.com/deepmind/optax) for gradient processing and
optimization, [chex](https://github.com/deepmind/chex) for reliable code and testing, and
In addition, DeepMind has open-sourced an [ecosystem of libraries around [Equinox](https://github.com/patrick-kidger/equinox) for neural networks.
JAX](https://deepmind.com/blog/article/using-jax-to-accelerate-our-research) (Watch the NeurIPS 2020 JAX Ecosystem at DeepMind talk
including [Optax](https://github.com/deepmind/optax) for gradient processing and [here](https://www.youtube.com/watch?v=iDxJxIyzSiM) for additional details.)
optimization, [RLax](https://github.com/deepmind/rlax) for RL algorithms, and
[chex](https://github.com/deepmind/chex) for reliable code and testing. (Watch
the NeurIPS 2020 JAX Ecosystem at DeepMind talk
[here](https://www.youtube.com/watch?v=iDxJxIyzSiM))
## Citing JAX ## Citing JAX