[D] Should We Be Using JAX in 2022?

1. flax

   1 11 7,233 9.7 Jupyter Notebook

   Flax is a neural network library for JAX that is designed for flexibility.

   

   What's your favorite Deep Learning API for JAX - Flax, Haiku, Elegy, something else?

2. SaaSHub

   www.saashub.com featured

   SaaSHub - Software Alternatives and Reviews. SaaSHub helps you find the best software and product alternatives

   

3. dm-haiku

   2 10 3,240 7.4 Python

   JAX-based neural network library

   

   What's your favorite Deep Learning API for JAX - Flax, Haiku, Elegy, something else?

4. elegy

   3 5 473 0.0 Python

   A High Level API for Deep Learning in JAX

   

   What's your favorite Deep Learning API for JAX - Flax, Haiku, Elegy, something else?

5. jax

   4 92 35,809 10.0 Python

   Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

   

   It really is quite hard to tell at this point. If we're talking just about Deep Learning, I think that JAX could be an awesome supplement to TensorFlow - since they both use XLA it's easy to move a model from JAX to TensorFlow, so hypothetically you could build in JAX and move to TF for deployment, but I don't know that that will be that useful in an industry setting.

6. equinox

   5 33 2,902 8.6 Python

   Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/

   

   Author of Equinox here. I'm glad to see it being mentioned in the wild!

7. diffrax

   6 23 2,047 7.7 Python

   Numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable. https://docs.kidger.site/diffrax/

   

   Some nice examples of this -- and in fact the whole reason Equinox exists -- can be found ubiquitously throughout Diffrax library. (A new JAX-based suite of diffeq solvers.) For example diffrax.AbstractSolver is an abstract parameterised function; diffrax.PIDController is a concrete instantiation of another abstract parameterised function. You can do some pretty cool stuff with this :)

8. jax-models

   7 6 161 0.0 Python

   Unofficial JAX implementations of deep learning research papers

   

   I've been using JAX, especially Flax for quite some time now for my reproducibility initiative (jax_models) and this is what I really appreciate about the framework

9. extending-jax

   8 2 403 3.0 Python

   Extending JAX with custom C++ and CUDA code

   

   You can check out this or this for more info. I think it is safe to assume that it is less stable than PyTorch - some other commenters have spoken about running into trouble with XLA in certain corner cases, but I have not experienced this so I can't speak to it.

Suggest a related project