Penzai

Bug fixes and improvements

- Conversion from Llama and Mistral models now correctly converts the RoPE wavelength (102)
- Conversion of all HuggingFace models now checks all changed attributes for possible incompatibilities (104)

Other changes

- Updates uv lockfile to use more recent package versions for development and tests.

New features and improvements

- Named array shape mismatch errors inside Linear layers now give a more informative error message based on the parameter name, if possible (92).
- The custom PyTree key `LayerStackGetAttrKey` used by `penzai.nn.layer_stack.LayerStack` no longer directly subclasses `jax.tree_util.GetAttrKey` (96), making it robust to future planned changes in JAX.

Other changes

- Project now includes a `uv` lockfile and can be easily installed using the `uv` package manager.
- Checks for formatting (via `pyink`), linting (via `pylint`), and type annotations (via `pytype`) are now run on every PR.
- `pyink` and `pylint` are configured to automatically run when editing the project in VSCode.
- Made minor tweaks to documentation configuration to avoid a build error.

- Added the ability to catch exceptions during rendering of `Selection` objects
- Added additional methods to Penzai variables to facilitate interoperability with other JAX ecosystem libraries
- Removed references to deprecated `named_shape` field of `jax.ShapeDtypeStruct`
- Added citation info for Penzai/Treescope paper (https://arxiv.org/abs/2408.00211)

New features:

- Added support for the Gemma 2 pretrained models.

Improvements:

- Adjusted colors for `TransformerBlock` and `Attention` when rendered with Treescope.
- Fixed an outdated link in the README.

Penzai 0.2.0 includes a number of substantial changes to the neural network and visualization systems, with the goal of improving usability and removing boilerplate.

⚠️ Breaking changes: ⚠️

- **`penzai.nn` now uses the V2 API.**
- The V2 neural network API, initially introduced in release 0.1.3 and previously available at `penzai.experimental.v2`, has been moved to the main Penzai namespace. References to `penzai.nn` or `from penzai import pz` will now refer to the V2 versions.
- The V2 API supports a number of additional features beyond V1:
- First-class mutable state and variable sharing
- Generalized transformer implementation with support for Llama, Mistral, and GPT-NeoX / Pythia
- Support for layer stacks and gradient checkpointing (to reduce memory usage)
- The original V1 API and data-effects system have been deprecated and moved to `penzai.deprecated.v1`.
- See the ["Changes in the V2 API"](https://penzai.readthedocs.io/en/v0.2.0/guides/v2_differences.html) for details and migration help.
- **Treescope has moved to a standalone package.**
- Penzai's interactive pretty-printer and array visualizer, Treescope, has moved into the separate [treescope package](https://github.com/google-deepmind/treescope).
- Treescope can be installed independently of Penzai, and now supports visualizing PyTorch models and tensors as well as JAX arrays and Penzai models. See the [Treescope 0.1.0 release notes](https://github.com/google-deepmind/treescope/releases/tag/v0.1.0) for details.
- Most top-level functions (e.g. `pz.ts.basic_interactive_setup` and `pz.show`) are still available through Penzai for compatibility, but users are encouraged to switch to using `treescope` directly instead.
- The extension API for custom node handlers has changed, and existing custom node handlers will need to be updated.

Other changes:

- Penzai no longer requires Equinox to be installed.
- `Selection.at(...)` now takes an explicit keyword argument `multiple` to distinguish single nodes from collections of nodes, instead of magically inferring this based on whether the collection existed in the tree.
- `Selection.at(...)` now correctly handles selecting the singleton nodes `None` and `()`.
- Outdated tutorials for the V1 API have been removed.

New features:

- Added a gradient checkpointing combinator `Checkpointed` in `penzai.experimental.v2.toolshed.gradient_checkpointing`. This can reduce memory usage when taking gradients through a model.

Bug fixes and improvements:

- Fixed an issue where rendering logic would run out of order in some situations (in particular, using a treescope rendering inside an iframe inside another treescope rendering).

Documentation changes:

- Added documentation of gradient checkpointing on the "how-to" guide page.