Penzai

- 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.

New features:

- Added `pz.nx.scan`, a named-axis wrapper around `jax.lax.scan` that handles tagging/untagging the scanned-over axis automatically.
- Improved support for dict-style indexing and index updates on `pz.nx.NamedArray` and `pz.nx.NamedArrayView`:
- Dict-style indexing of named arrays now supports integer advanced indexing with positional axes.
- Dict-style index updates (e.g. `array.at[{ "foo": 3 }].set( ... )`) are now supported.

Bug fixes and improvements:

- Output views in JupyterLab (and other non-sandboxed IPython notebook interfaces) should no longer interfere with each other.
- Creating multiple output views no longer corrupts the rendering.
- Saved outputs from different Penzai versions no longer interfere due to global browser Javascript state.
- Added CSS scoping and content visibility annotations to Treescope renderings, which should improve performance in modern browsers for documents with many treescope renderings.

Documentation changes:

- Fixed broken links in the README and documentation homepage.
- Updated the image in the README.
- Added some details on how to customize the attention masks in the "how-to" guide page.

New features:

- **V2 neural network API** (`penzai.experimental.v2`)
- A redesign of Penzai's neural network system, which introduces first-class mutable state and variable sharing, and removes boilerplate.
- You can read more about the differences and how to migrate [here](https://penzai.readthedocs.io/en/v0.1.3/guides/v2_differences.html).
- We plan to replace the original neural network system with this V2 API in Penzai release 0.2.0.
- **Llama, Mistral, and GPT-NeoX / Pythia support**
- The pretrained transformer implementation has been generalized, and now supports Llama, Mistral, and GPT-NeoX / Pythia pretrained models.
- (This implementation is specific to the V2 neural network API.)
- Other features:
- New `LayerStack` combinator, which uses `jax.lax.scan` to efficiently repeat layers with the same structure
- Named arrays can now be updated using `.at[...].set(...)` operations. For now, only positional indexing is supported (with broadcasting over named axes).

Bug fixes and improvements:

- Fixed issue where unit test discovery was not picking up tests in subdirectories (38)
- Fixed issue where adding a NamedArray to a JAX array would not correctly lift the JAX array to a NamedArray (37)

Documentation changes:

- Added documentation of the V2 API, along with instructions on how to migrate.
- Added a "How-To Guide" for common tasks (V2 API only)