Remove TF dependency in Lookahead Optimizer on MNIST example

[rajasekharporeddy]

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

[vroulet]

Looks great thanks

[rajasekharporeddy]

Hi @vroulet, I forgot to remove the .as_numpy_iterator() while enumerating over the batched data loader in the previous commit. Please review the changes once more.

[rdyro]

I think you're running into some doctests issues, can you take a look at them?

You should be able to test locally by doing

$ cd optax
$ uv pip install ".[docs]"
$ cd docs
$ make test

[rdyro]

Testing this locally, you just need to rebase to HEAD and it should start passing the tests.