Add more examples in the NNX docs

[vfdev-5]

We plan for moving Jax AI Stack examples to NNX docs:

• (Low prio) Part 2: Debug a variational autoencoder (VAE)
• Part 3: Train a diffusion model for image generation (@samanklesaria, Add diffusion example ported from Jax AI stack #5403)
• Visualize JAX model metrics with TensorBoard (@samanklesaria, Update mnist tutorial to use tensorboard #5425)
• Introduction to Data Loaders on CPU with JAX / Introduction to Data Loaders on GPU with JAX (@samanklesaria Add data loaders tutorial #5454 )
• JAX for PyTorch users / Porting a PyTorch model to JAX (@vfdev-5, [Docs] migration guide "From PyTorch to JAX and Flax"  #5408)
• Train a miniGPT language model with JAX (@samanklesaria, Port Jax AI stack minigpt example #5405)
• Text classification with a transformer language model using JAX
• Machine Translation with encoder-decoder transformer model (@samanklesaria, Add machine translation example #5431)
• Image segmentation with UNETR model
• Image Captioning with Vision Transformer (ViT) model (@vfdev-5, ...)
• Train a Vision Transformer (ViT) for image classification with JAX (@vfdev-5, Added ViT training example #5455)
• Time series classification with CNN