Skip to content

Add diffusion example ported from Jax AI stack - #5403

Merged
copybara-service[bot] merged 2 commits into
google:mainfrom
samanklesaria:diffusion
Apr 13, 2026
Merged

copybara-service[bot] merged 2 commits into
google:mainfrom
samanklesaria:diffusion

Conversation

@samanklesaria

Copy link
Copy Markdown
Collaborator

This ports the example from the old Jax AI stack website to modern NNX. I switched out flow matching for DDPMs.

@review-notebook-app

Copy link
Copy Markdown

Check out this pull request on  ReviewNB

See visual diffs & provide feedback on Jupyter Notebooks.


Powered by ReviewNB

@samanklesaria
samanklesaria marked this pull request as draft April 10, 2026 14:52
@samanklesaria
samanklesaria force-pushed the diffusion branch 2 times, most recently from 8f6c218 to de86260 Compare April 10, 2026 15:51
@samanklesaria
samanklesaria marked this pull request as ready for review April 10, 2026 16:00
@copybara-service
copybara-service Bot merged commit 9e73be4 into google:main Apr 13, 2026
20 checks passed
Transformer encoder trained on the One Billion Word Benchmark.


- :octicon:`mark-github;0.9em` `Diffusion Models <https://github.com/google/flax/tree/main/examples/digits_diffusion_model/>`__ :

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@samanklesaria this link is unvalid, we should expose the notebook instead?

@samanklesaria samanklesaria mentioned this pull request Apr 20, 2026
9 of 14 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants