Skip to content

Update mnist tutorial to use tensorboard - #5425

Merged
copybara-service[bot] merged 2 commits into
google:mainfrom
samanklesaria:tensorboard_example
May 12, 2026
Merged

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

Conversation

@samanklesaria

@samanklesaria samanklesaria commented Apr 17, 2026 •

Copy link
Copy Markdown
Collaborator

This PR ports the example https://docs.jaxstack.ai/en/latest/JAX_visualizing_models_metrics.html from the Jax AI stack. It also switches the import of MNIST to huggingface datasets rather than tfd. For real world examples, we use grain over tfd anyway, so teaching users about using tfd seems unnecessary. We can also avoid the heavy weight tensorflow dependency this way.

@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 requested a review from vfdev-5 April 17, 2026 02:41
@samanklesaria
samanklesaria marked this pull request as ready for review April 17, 2026 15:21
Comment thread docs_nnx/mnist_tutorial.ipynb
from IPython.display import clear_output
import matplotlib.pyplot as plt
from tensorboardX import SummaryWriter
import tensorboard

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.

Why do we need to import tensorboard here?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

For the '%tensorboard' magic to visualize within Jupyter. Unfortunately, the within-notebook Jupyter doesn't seem to show up in the Sphinx rendering. Perhaps it's easiest no to use the within-notebook tensorboard magic and instead put in screetshots like the original jax AI stack one does.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

I added a screenshot (because the %jupyter magic doesn't survive Sphinx), but kept the use of the magic.

@vfdev-5 vfdev-5 mentioned this pull request Apr 20, 2026
9 of 14 tasks
Comment thread docs_nnx/mnist_tutorial.md Outdated
@copybara-service
copybara-service Bot merged commit a9057b8 into google:main May 12, 2026
21 checks passed
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