jaxviz - Package to interactively visualize Jax models in notebooks

[sachinhosmani]

I have built a new open source tool called jaxviz to interactively visualize the forward pass of Jax models in notebooks.

It produces a hierarchical graph view on a canvas which you can drag, zoom, expand and collapse.

It also shows the compiled per-device view for sharded models, which includes collective operations like AllReduce.

You can check out some examples on the website here, or just pip install jaxviz in your notebook to try it out on a real model.

Keen to hear your feedback on it!

Thank you