PR #5077 Migrate VAE example to Flax NNX

Migrate VAE example to Flax NNX #5077

sanepunk wants to merge 7 commits into google:main from sanepunk:main

Migrated from nn -> nnx.

b4627cbc

vfdev-5 commented on 2025-11-07

Update examples/vae/main.py

6fd002a4

Necessary changes made as per the comments by @vfdev-5.

f959a6de

Removed unused packages ml-collections and clu

76971ef6

pre-commit changes

ed731fc0

vfdev-5 changed the title ~~Migrate VAE example to Flax NNX with JIT optimization - 9.2x speedup~~ Migrate VAE example to Flax NNX 86 days ago

ghost commented on 2025-11-10

vfdev-5 commented on 2025-11-10

Refactor VAE example to use dataclass for configuration and simplify …

f9ede36e

Merge branch 'main' into main

07a945ce

sanepunk requested a review from

vfdev-5 58 days ago

Reviewers

vfdev-5

Assignees

No one assigned

Labels

None yet

Milestone

No milestone