Penzai
A JAX research toolkit to build, edit, & visualize neural networks
...It is designed to make machine learning research more interpretable and interactive, particularly for tasks like model surgery, ablation studies, architecture debugging, and interpretability research. Unlike conventional neural network libraries, Penzai exposes the full internal structure of models, enabling fine-grained inspection and modification after training. Its modular design includes tools for tree manipulation, named axes, and declarative neural network construction. The library integrates tightly with Treescope, an advanced pretty-printer for visualizing deeply nested JAX pytrees and NDArray structures. ...