Is Google replacing TensorFlow with JAX

TensorFlow is a popular deep learning framework that has been used by researchers and developers for many years. However, there is a new framework on the rise called JAX. JAX is a high-performance numerical computation library that is built on top of NumPy. It offers a number of advantages over TensorFlow, including:

  • Speed: JAX is significantly faster than TensorFlow for many tasks.
  • Ease of use: JAX is easier to use than TensorFlow for some tasks, especially for those who are familiar with functional programming.
  • Flexibility: JAX is more flexible than TensorFlow, allowing you to do things that are not possible with TensorFlow.

So, with all of these advantages, is it possible that Google will replace TensorFlow with JAX? It is certainly possible, but it is not clear that it will happen anytime soon. TensorFlow has a large user base and a large ecosystem of libraries and tools. It would be a major undertaking to replace TensorFlow with JAX.

However, Google is clearly investing in JAX. The TensorFlow team has been working on integrating JAX into TensorFlow, and there are a number of Google researchers who are using JAX for their research. It is possible that Google will eventually make JAX the default deep learning framework for Google’s product?

JAX is significantly faster than TensorFlow for many tasks

Since JAX is just a computation library running in Python, you are probably comparing the performance with it and using TensorFlow from Python as well. Here, it is possible that JAX is faster and simpler to use. Though, you still need a Python process to run it.

I think where TensorFlow still wins is that you can export your trained graph to a binary that can then be loaded and served by a non-Python process, like TFServing or Java. I think that is where the comparison should be done: how fast JAX can deliver a live prediction compared to TensorFlow, taking into account the different runtimes they can run on.

Thanks for bringing up the question, I am quite curious myself. Can you elaborate on the points? I don’t think that any of them is true, but I am glad to learn more.

Speed: is it? TensorFlow can also do XLA compilation (with tf.function(jit=True), so they end up in a similar layer. There can still be optimizations that differ, but are there? Glad to learn about any comparison.
Ease of use: that surely depends on what you want to do, JAX for example misses stateful variables that can make life easier, but most importantly, TensorFlow also contains all the functions that JAX has, AFAIU. The tf.math. functions.
Flexibility: is it? Can you give a single example maybe? Again, I am glad to learn about that as I can’t think of one.

Tensorflow is pretty dead. Just look at the activity of this forum. It would not be surprising if Google abandoned it.