Excited to share Penzai, a JAX research toolkit from @GoogleDeepMind for building, editing, and visualizing neural networks! Penzai makes it easy to see model internals and lets you inject custom logic anywhere. Check it out on GitHub: github.com/google-deepmin…
Penzai integrates seamlessly with @GoogleColab and the JAX ecosystem. It represents models as legible, editable data structures, to help researchers understand and modify them after they are trained. Built with support from @DougalMaclaurin, @dtarlow2, and @hugo_larochelle!
@_ddjohnson @GoogleDeepMind nice!! how long should it take to load the visualizations in colab? im trying the How to Think in Penzai notebook and all the cells just look like this: using L4 GPU
@_ddjohnson @GoogleDeepMind Looks really cool! Qq, can you port existing flax models weights to this, without the initial model being written in penzai.nn?
@_ddjohnson @GoogleDeepMind Dang, I need more hours in the day to check this out. Amazing work.