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!