Treescope

Treescope is an interactive HTML pretty-printer and N-dimensional array ("tensor") visualizer, designed for machine learning and neural networks research in IPython notebooks. It serves as a drop-in replacement for the standard IPython/Colab renderer, enhancing output with features like expandable/collapsible subtrees, embedded faceted visualizations of arbitrary-dimensional arrays, and color-coding for model structures. It supports inspecting data from libraries such as JAX, NumPy, PyTorch, Equinox, Flax NNX, and Penzai. Maintained by Google DeepMind, the current version is 0.1.10.