G github.com Penzai: JAX research toolkit for building, editing, and visualizing neural nets 261 points mccoyb 2 years ago 51 comments