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.

[newresu]

@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

[openSourcerer9000]

Why would we think JAX would have a longer shelf life then?

[pass-lin]

Does Google have no one to manage the tf community now, and the issues are all unknown message. It's crazy

[rivershah]

I have moved to backend-agnostic keras and gathered fresh benchmarks.

jax is no panacea. it still carries multiple gotchas and runs slower on l4, h100, tpu-v5p and tpu-v6e across cnns, transformers, and gnns. Compile time, training throughput, and inference latency mostly favour tensorflow. There is no consistent winner.

keras + tensorflow + dtensor remains the speed leader for these workloads. google has withdrawn nearly all high-level tensorflow-on-tpu guides; only the bare tpu-vm compatibility tables remain. Dear google we may not be $100M a year gpu/ tpu customers but six-figure-spend gcp customers still require full tensorflow coverage.

tensorflow needs more engineering investment, not less. Promote the experimental data-pipeline and kernel modules that give even more speedups, excise obsolete code even at the cost of minor breakage, and keep advancing kernel fusion. Pulling support and documentation is the wrong move.

[rivershah]

I take much of what I said back...over the last two weeks I ground away on my jax based models and now getting extremely good performance and code ergonomics on larger tpus.

When I wrote the above message, I was in the trough of sadness of having to say goodbye to tensorflow but now I am elated to see what I can do with keras + jax.