Has Google completely abandoned TensorFlow?

[edwardyehuang]

Starting this year, I have raised some issues related to TensorFlow, but no one has addressed them at all. Recently, I have also submitted some pull requests to the TensorFlow repository, but no one has reviewed them. Does this mean that Google has abandoned TensorFlow? If so, I will shift my development focus to Jax.

[abhaskumarsinha]

Comments: I don't think Google has completely abandoned TF. But their main focus seems to be on JAX. Currently, in my opinion, JAX is the fastest, easiest, and simplest backend to run deep learning models on.
A Google's official engineer might comment better here.
But I often see this trend, in other repos, including - XLA (Accelerated Linear Algebra Compiler and so on).

[abhaskumarsinha]

@ghsanti They do a lot of things imo, but not a single thing properly. As @edwardyehuang pointed out, even PR reviews take like months to years for them to acknowledge. With such pace, they'll never bring up a thing of quality.

[edwardyehuang]

@fchollet @qlzh727 @jeffcarp @SamanehSaadat Any comments? As a researcher with limited time, I want to focus my limited time on a framework that provides support.

[abhaskumarsinha]

@edwardyehuang What exactly are you working on? I may suggest a better alternative if it is in my scope of knowledge?

[edwardyehuang]

@edwardyehuang What exactly are you working on? I may suggest a better alternative if it is in my scope of knowledge?

Thanks, I'm fine, and I'm also very familiar with PyTorch and Jax. I'm just considering whether it's time to migrate some of my core code out of TensorFlow.

[abhaskumarsinha]

@edwardyehuang I'm not an expert, but migrating to Keras would be worth, you might probably get support of PyT, JAX and TF simultaneously if possible.

[fchollet]

TensorFlow is still maintained, and will keep being maintained -- at least in the medium term. Most production ML systems at Google use TensorFlow and getting them all to migrate would be a multi-year project.

However JAX is definitely seeing more development capacity at this time. IMO it's a good idea to switch to JAX, not so much because of fear that TF might not be supported in the future, but simply because JAX is a technologically better platform to develop on. In particular it is easier to scale (via sharding APIs, e.g. keras.distribution APIs) and it is generally more performant.

My personal style is to keep my code backend-agnostic and to switch from one backend to another based on current needs (might use torch to debug, JAX or TF to run). I still use many TF components though, in particular tf.data is indispensable.

[ghsanti]

@fchollet

use many TF components though, in particular tf.data is indispensable.

Do you prefer it over PyDataset for performant training?

[rkazants]

@fchollet, I see some advantage to use TensorFlow for language models. Namely, I can incorporate string tensor pre-processing part into a model and get e2e language models. PyTorch and JAX do not have string tensor support so I have to keep tokenizers out-of-the model that is not convenient for production and requires additional dependencies, memory copies in run-time.
How do you think? Is it important problem for PyTorch and JAX?

[fchollet]

Do you prefer it over PyDataset for performant training?

When you pass a PyDataset to fit() it actually gets converted to a tf.data.Dataset internally (when training with TF).

In general yes you should prefer tf.data over PyDataset when feasible.

I see some advantage to use TensorFlow for language models. Namely, I can incorporate string tensor pre-processing part into a model and get e2e language models.

Is your concern with training in a Python based runtime, or with serving in non-Python runtime?

For serving, keeping preprocessing inside the TF graph is super nice for sure. For training it doesn't make as big a difference -- you can just map the preprocessing in your tf.data pipeline. KerasNLP task models actually do this for you automatically (so they "feel" end-to-end, regardless of the backend you use).

One thing you could do is training with the JAX backend, and switch to the TF backend with a preprocessing layer or tokenizer as part of the model before you export a TF SavedModel for serving.

[innat-asj]

@edwardyehuang

As a researcher with limited time, I want to focus my limited time on a framework that provides support.

Just move to PyTorch.

[edwardyehuang]

Yes, PyTorch + XLA seems one of the good choices for me