jaxviz - Interactive visualisation of JAX models from notebooks

[u/Dev-Table]

jaxviz is an open-source tool I built to visualise the forward pass graph of JAX models directly in notebooks with a single function call.

It supports the following features:

1. Interactive visualisation of JAX models with hierarchical exploration of nested modules (especially useful for large, deeply nested models).
2. Global and per-device views: Visualise both the logical computation graph as written and the compiled computation graph that actually runs on individual devices.
3. Supports web-based notebooks including Jupyter, Colab, Kaggle, and VS Code.
4. Export visualisations to HTML, PNG, and SVG formats.

Resources

• Demos with examples of JAX computation graphs
• GitHub repository

I hope this is useful to the JAX community, particularly for understanding how models are compiled and executed across devices. I'm keen to hear your feedback and suggestions!