# google-deepmind/penzai

A JAX research toolkit for building, editing, and visualizing neural networks.

Repository: https://github.com/google-deepmind/penzai
Canonical: https://ross.abutalabs.com/products/penzai
Homepage: https://penzai.readthedocs.io/
Language: Python
License: Apache-2.0
License Family: permissive
Topics: fine-tuning, interpretability, jax, neural-networks, visualization
Archived: true
Last push: 2025-06-22T05:45:13+00:00

## Health v2 (maintenance only)
Score: 10/100 (v2, computed 2026-09-03T02:20:16.233290+00:00)
- activity 28, release rhythm 28, longevity 63
- inputs: {"age_days": 882, "days_push": 437, "days_rel": 512, "gap_med": 71, "n_releases_24m": 4}
- flags: archived
- formula: round(0.45*activity + 0.35*rhythm + 0.20*longevity); archived -> min(score, 10)

## Adoption (not part of the score)
Stars 1901, forks 70 (observed 2026-08-28T04:05:51.649951+00:00)

## What it is
Penzai is a JAX research toolkit for building, editing, and visualizing neural networks as legible, functional pytree data structures. It includes tools for model surgery, interpretability research, activation inspection, and interactive visualization of deeply nested model structures.

## Use cases
- reverse-engineer and ablate components of trained neural networks
- inspect and probe internal activations of transformer models
- perform model surgery on pretrained Gemma, Llama, Mistral, or GPT-NeoX models
- visualize deeply nested JAX pytrees and NDArrays in notebooks
- build and train custom neural network architectures declaratively
- debug neural network architectures interactively
- fine-tune large language models with editable structures

## When to choose
- you need to modify or analyze trained models rather than just train them
- you want interpretability research tooling for transformer LLMs in JAX
- you prefer models as inspectable, editable pytree data structures
- you need interactive visualization of model internals in IPython or Colab

## When to avoid
- you need a production training framework with mature ecosystem like PyTorch
- you want a simple high-level API like Keras for standard workflows
- your project is not based on JAX or the Python ML ecosystem
- you need distributed large-scale training infrastructure out of the box

## Facets
- artifact type: library
- maturity: active
- function: machine-learning, deep-learning, data-visualization, developer-tools
- domain: machine-learning, deep-learning, large-language-models, developer-tools
- platform: python, jvm
- tags: jax, interpretability, model-surgery, transformers, pytrees, neural-networks, fine-tuning, research-toolkit

## Member repositories
- google-deepmind/penzai (main) score 10

## Provenance
- Observed fields: from GitHub, fetched 2026-08-28T04:05:51.649951+00:00.
- Health v2: computed from the inputs above; adoption is never an input.
- Inferred fields (summary, facets, guidance): AI-extracted, prompt v1, taxonomy v1, on 2026-08-30T03:11:56.543876+00:00, confidence not recorded.
  - readme: https://github.com/google-deepmind/penzai (fetched 2026-08-28T04:05:51.649951+00:00, sha 0a65114c31e9)
  - registry_pypi: https://pypi.org/pypi/penzai/json (fetched 2026-08-29T10:51:35.332162+00:00, sha e9df2f652c40)
- Data as of 2026-08-30T08:39:29.467469+00:00.