Search code, repositories, users, issues, pull requests...

submited by

Style Pass

2024-04-19 16:00:05

盆 ("pen", tray) 栽 ("zai", planting) - an ancient Chinese art of forming trees and landscapes in miniature, also called penjing and an ancestor of the Japanese art of bonsai.

Penzai is a JAX library for writing models as legible, functional pytree data structures, along with tools for visualizing, modifying, and analyzing them. Penzai focuses on making it easy to do stuff with models after they have been trained, making it a great choice for research involving reverse-engineering or ablating model components, inspecting and probing internal activations, performing model surgery, debugging architectures, and more. (But if you just want to build and train a model, you can do that too!)

penzai.nn (pz.nn): A declarative combinator-based neural network library and an alternative to other neural network libraries like Flax, Haiku, Keras, or Equinox, which exposes the full structure of your model's forward pass in the model pytree. This means you can see everything your model does by pretty printing it, and inject new runtime logic with jax.tree_util. Like Equinox, there's no magic: models are just callable pytrees under the hood.

penzai.treescope (pz.ts): A superpowered interactive Python pretty-printer, which works as a drop-in replacement for the ordinary IPython/Colab renderer. It's designed to help understand Penzai models and other deeply-nested JAX pytrees, with built-in support for visualizing arbitrary-dimensional NDArrays.