flax
Add diffusion example ported from Jax AI stack
#5403

Merged

Add diffusion example ported from Jax AI stack #5403

copybara-service merged 2 commits into google:main from samanklesaria:diffusion

Add diffusion example

e8a62a5e

samanklesaria marked this pull request as draft 86 days ago

samanklesaria force pushed from 0b42300f to 8f6c2183 86 days ago

samanklesaria force pushed from 8f6c2183 to de862604 86 days ago

Use flow matching for diffusion example

0f1ef741

samanklesaria force pushed from de862604 to 0f1ef741 86 days ago

samanklesaria marked this pull request as ready for review 86 days ago

IvyZX approved these changes on 2026-04-13

IvyZX added pull ready

copybara-service merged 9e73be40 into main 83 days ago

vfdev-5 commented on 2026-04-16

Reviewers

IvyZX

vfdev-5

Assignees

No one assigned

Labels

pull ready

Milestone

No milestone