remove tensorflow as dependency for the examples

[fabianp]

Currently we're importing tensorflow for some examples such as in https://optax.readthedocs.io/en/latest/_collections/examples/cifar10_resnet.html.

It should be possible to avoid this import and instead rely exclusively on tfds and (if necessary grain)

See for example https://www.tensorflow.org/datasets/tfless_tfds on how to import tfds datasets without tensorflow.

[mmhamdy]

Anyone working on this? if not, I'd be happy to work on it.

[fabianp]

All yours @mmhamdy !

[fabianp]

hey @mmhamdy , any progress on this? I might have some cycles to work on this issue next week if it's not yet solved

[mmhamdy]

Hi @fabianp, sorry for taking that long, January was such a MONTHFUL😅 for me! But anyway, I'll submit a PR in a couple of days.

I was leaving grain as a last resort and trying to use TFDS only. The thing is, the promise of a TFless TFDS in the documentation isn't a complete one, quoting the doc:

In future versions, we are also going to make the dataset preparation TensorFlow-free.

So, we are getting a TF-free dataset but we still need TF to process it. That's nice, but still not that helpful.

I used grain in the end, as you might've expected. It's working fine, just a small problem with MNIST. When loaded using the ArrayRecord-based data source, we get a chunk header hash mismatch and can't access it. I've opened an issue and still working on it.

The other datasets are fine. Now, the examples notebooks that use TF and TFDS are:

adversarial_training.ipynb
lookahead_mnist.ipynb
mlp_mnist.ipynb
differentially_private_sgd.ipynb
cifar10_resnet.ipynb
reduce_on_plateau.ipynb

The first 4 examples use MNIST, and cifar10_resnet.ipynb still uses TF for image augmentation. Also, a small note, mlp_mnist.ipynb and differentially_private_sgd.ipynb appear on the gallery but don't appear on the side menu.

So, I think I'm going to start with reduce_on_plateau.ipynb. One last question, do you prefer a separate PR for each notebook or just one PR for all?

Thank you, and sorry for the long reply. Was just thinking out loud.

[fabianp]

Excellent, and thanks for the update! I'm trying to push for this as I'd like to test on Python 3.12 but unfortunately TF is blocking us from doing so 😞 (https://github.com/google-deepmind/optax/pull/714)

So, we are getting a TF-free dataset but we still need TF to process it. That's nice, but still not that helpful.

That's still an improvement 😄 . Let's take it step by step and merge this part, then we can take care of the processing in a separate PR :-)

I used grain in the end, as you might've expected. It's working fine, just a small problem with MNIST. When loaded using the ArrayRecord-based data source, we get a chunk header hash mismatch and can't access it. I've opened an issue and still working on it.

huh, that's annoying. Could you link this issue here? perhaps I can take a look into it or ask around.

The other datasets are fine. Now, the examples notebooks that use TF and TFDS are:

adversarial_training.ipynb
lookahead_mnist.ipynb
mlp_mnist.ipynb
differentially_private_sgd.ipynb
cifar10_resnet.ipynb
reduce_on_plateau.ipynb

The first 4 examples use MNIST, and cifar10_resnet.ipynb still uses TF for image augmentation. Also, a small note, mlp_mnist.ipynb and differentially_private_sgd.ipynb appear on the gallery but don't appear on the side menu.

So, I think I'm going to start with reduce_on_plateau.ipynb. One last question, do you prefer a separate PR for each notebook or just one PR for all?

yeah, let's start with one and take it from there. It's always easier to review a small PR and iterate on it rather than making huge changes to the codebase on one go.

Thank you, and sorry for the long reply. Was just thinking out loud.

Thanks to you! I think the tag "good first issue" was over-optimistic on my end 😅 , this is really a quite difficult issue!

[fabianp]

Also, a small note, mlp_mnist.ipynb and differentially_private_sgd.ipynb appear on the gallery but don't appear on the side menu.

Nice catch, I've opened https://github.com/google-deepmind/optax/issues/767 to track this

[mmhamdy]

huh, that's annoying. Could you link this issue here? perhaps I can take a look into it or ask around.

Yeah, but maybe there is still little hope here. This is the issue

yeah, let's start with one and take it from there. It's always easier to review a small PR and iterate on it rather than making huge changes to the codebase on one go.

Sure, I'll finish the fashion-mnist example and submit it in a PR. It is working nicely with grain, although a tiny bit slower, I think. But I'm still discovering grain and maybe could make it faster. If everything goes smoothly without any surprises, this will be the first truly TF-free notebook and also the first to feature grain loader😃

I think the tag "good first issue" was over-optimistic on my end 😅 , this is really a quite difficult issue!

Yeah, I agree (just remembered I intended to do it while migrating the jaxopt example 😅😅). But, to be honest, I'm really enjoying it.