[Docs] JAX_Vision_transformer: ViT Tutorial at https://docs.jaxstack.ai/en/latest/JAX_Vision_transformer.html Fails on Kaggle (TPU v5e-8) or T4 (Colab)

# ViT Tutorial at https://docs.jaxstack.ai/en/latest/JAX_Vision_transformer.html fails on Kaggle (TPU v5e-8)

I noticed it uses older NNX APIs such as nnx.ModelAndOptimizer, which are no longer present in the current Flax NNX version, so the tutorial does not run on TPU v5e-8 (Kaggle) or T4 (Colab).

## Updating https://docs.jaxstack.ai/en/latest/JAX_Vision_transformer.html
I have made a Jax-Flax ViT tutorial Repository at: https://github.com/Prakharprasun/jax-flax-vit-pipeline/
It avoids the old NNX APIs and runs correctly on TPU v5e-8.

I'd be happy to help by:
1. Updating the existing tutorial using modern Flax patterns or
2. Contributing a new ViT fine-tuning tutorial based on the updated approach

Whichever direction the maintainers prefer.
I can work on this over the coming weeks and will follow any guidelines or structure you recommend.

Thank you!

Start by running the ViT tutorial at https://docs.jaxstack.ai/en/latest/JAX_Vision_transformer.html on Kaggle TPU v5e-8 or Colab T4 and identify the outdated Flax NNX APIs, including nnx.ModelAndOptimizer. Compare it with the modern approach in https://github.com/Prakharprasun/jax-flax-vit-pipeline/. Done means the existing tutorial, or an agreed replacement tutorial, runs successfully on the reported environments.

Written by the indexing model from the issue text.