GitHub repository: mahdi-shafiei/penzai
Primary language: Python
Description: A JAX research toolkit for building, editing, and visualizing neural networks.