Migrate VAE example to Flax NNX#5077

Open

sanepunk wants to merge 5 commits intogoogle:mainfrom

Commits on Nov 7, 2025

Migrated from nn -> nnx.
sanepunk
committed

Commits on Nov 8, 2025

Commits on Nov 9, 2025