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)