The MNIST training example could be better

[bjourne]

The notebook on training on MNIST with JAX (https://docs.jax.dev/en/latest/notebooks/neural_network_with_tfds_data.html) has several problems:

• It uses tanh activations.
• Scale and learning rate is set to 1e-2, which leads to very slow convergence compared to 1-e1.
• It uses tensorflow-datasets, which does not support Python 3.14. And since JAX replaces TensorFlow the dependency on an (outdated?) framework is odd.
• The code is quite different from mnist_classifier.py and mnist_classifier_fromscratch.py so it's difficult to know whether it reflects current best practices or not. In particular, random numbers are handled differently in the scripts in the examples directory.