[D] Should We Be Using JAX in 2022?

1. flax

   1 10 6,677 9.5 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. Stream

   getstream.io featured

   Stream - Scalable APIs for Chat, Feeds, Moderation, & Video. Stream helps developers build engaging apps that scale to millions with performant and flexible Chat, Feeds, Moderation, and Video APIs and SDKs powered by a global edge network and enterprise-grade infrastructure.

   

3. dm-haiku

   2 10 3,057 6.9 Python

   JAX-based neural network library

   

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

4. elegy

   3 5 475 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 90 32,824 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,428 8.8 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 1,695 8.4 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 156 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. InfluxDB

   www.influxdata.com featured

   InfluxDB – Built for High-Performance Time Series Workloads. InfluxDB 3 OSS is now GA. Transform, enrich, and act on time series data directly in the database. Automate critical tasks and eliminate the need to move data externally. Download now.

   

10. extending-jax

    8 2 398 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