Skip to content

Update mnist tutorial to use tensorboard#5425

Open
samanklesaria wants to merge 2 commits intogoogle:mainfrom
samanklesaria:tensorboard_example
Open

Update mnist tutorial to use tensorboard#5425
samanklesaria wants to merge 2 commits intogoogle:mainfrom
samanklesaria:tensorboard_example

Conversation

@samanklesaria
Copy link
Copy Markdown
Collaborator

@samanklesaria samanklesaria commented Apr 17, 2026

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.

Comment thread docs_nnx/mnist_tutorial.md Outdated
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants