Remove TF dependency in reduce_on_plateau example

[mmhamdy]

This PR removes Tensorflow as a dependency in reduce_on_plateau.ipynb, and uses only TFDS, with grain for loading and preparing the dataset.

Addressing issue #656

[fabianp]

This looks amazing! Way to go @mmhamdy! 🚀

Just one thing, could you please remove the %pip install statements, and instead add these dependencies to the pyproject.toml file, under section [docs].

[mmhamdy]

Sure, much more cleaner this way.

[mmhamdy]

Hey @fabianp, looks like there is a problem with python 3.9 and grain.

[fabianp]

ouch! yeah, I'm seeing from https://github.com/google/grain/blob/main/pyproject.toml that they require Python 3.10 ....

[mmhamdy]

ouch indeed! 😅 What do you think? should we stick to notebook downloads until we get rid of either TensorFlow or Python 3.9?

[fabianp]

I'm unsure. Let me talk today with the other optax devs to see what they think

…

On Sun, Feb 4, 2024, 19:50 Mohammed Hamdy ***@***.***> wrote: ouch indeed! 😅 What do you think? should we stick to notebook downloads until we get rid of either TensorFlow or Python 3.9? — Reply to this email directly, view it on GitHub <#768 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AACDZBYXEP4G7ZGONXUNHQDYR7Q6XAVCNFSM6AAAAABCYTNHXOVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTSMRVHA4TKMZRGE> . You are receiving this because you were mentioned.Message ID: ***@***.***>

[fabianp]

@mmhamdy : I talked today with the pygrain devs, we're going to try to make it work on Python 3.9 . Give me a couple of days to try to get this working 😉

[mmhamdy]

Sounds great! Also, the MNIST issue has been solved, so I'll be working on the other examples in the meantime.

[fabianp]

BTW I submitted a PR fixing python 3.9 errors: google/grain#338

[mmhamdy]

That's QUITE a PR! 😅😅

[fabianp]

yeah, one thing that worries me a bit more about pygrain is that it it doesn't seem to build on OSX (tried and failed), and probably also Windows (haven't tried there). Before committing replacing TF with pygrain it would be important to make sure that those platforms are supported (just to say that it might take a bit before we merge these PRs unfortunately)

[mmhamdy]

Well, replacing TensorFlow wasn't expected to be a walk in the park anyway 😃. But we've come a long way and I see that the people at pygrain are working daily on it.

[fabianp]

thanks for updating @mmhamdy ! Very glad to see that the test now pass!

[vroulet]

For the record, in #1270 we removed this notebook (and others) from the docs building configuration so we don't depend on tensorflow anymore during the continuous integration.
This PR is still great to actually migrate the examples to grain for the sake of users learning how to use grain etc... And by the way, grain now supports mac os (not windows I believe but that's fine). So we could try merging this @mmhamdy if you want

[mmhamdy]

Hi @vroulet, yeah sure that would be great! It has been some time! Let's close this PR and I'll open a new one.