[D] Should We Be Using JAX in 2022?

This page summarizes the projects mentioned and recommended in the original post on /r/MachineLearning

Scout Monitoring - Free Django app performance insights with Scout Monitoring
Get Scout setup in minutes, and let us sweat the small stuff. A couple lines in settings.py is all you need to start monitoring your apps. Sign up for our free tier today.
www.scoutapm.com
featured
InfluxDB - Power Real-Time Data Analytics at Scale
Get real-time insights from all types of time series data with InfluxDB. Ingest, query, and analyze billions of data points in real-time with unbounded cardinality.
www.influxdata.com
featured
  • flax

    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?

  • Scout Monitoring

    Free Django app performance insights with Scout Monitoring. Get Scout setup in minutes, and let us sweat the small stuff. A couple lines in settings.py is all you need to start monitoring your apps. Sign up for our free tier today.

    Scout Monitoring logo
  • dm-haiku

    JAX-based neural network library

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

  • elegy

    A High Level API for Deep Learning in JAX

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

  • jax

    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.

  • equinox

    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!

  • diffrax

    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 :)

  • jax-models

    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

  • InfluxDB

    Power Real-Time Data Analytics at Scale. Get real-time insights from all types of time series data with InfluxDB. Ingest, query, and analyze billions of data points in real-time with unbounded cardinality.

    InfluxDB logo
  • extending-jax

    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.

NOTE: The number of mentions on this list indicates mentions on common posts plus user suggested alternatives. Hence, a higher number means a more popular project.

Suggest a related project

Related posts