Penzai

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)

New features:
- Penzai's pretty printer Treescope can now be more easily extended to work with custom types, by defining the special method `__penzai_repr__`, which will be detected during pretty-printing. Renderings of these custom types can be built using `penzai.treescope.repr_lib`, a stable extension API that can be used to implement `__penzai_repr__` in a forwards-compatible way. (Note that the exact internal representation of Treescope renderings is still unstable and may change in the future; see [the documentation for `penzai.treescope`](https://penzai.readthedocs.io/en/v0.1.2/api/treescope.html#adding-treescope-support-for-new-types) for details.)
- Added `pz.ts.basic_interactive_setup()`, which simplifies setting up the pretty-printer in an IPython notebook.

Bug fixes and improvements:
- Improved detection of mutable vs immutable objects during pretty-printing.
- Fixed detection of special methods `_repr_html_` and `__penzai_repr__` to avoid objects that generate arbitrary attributes dynamically.
- Copy-path buttons are no longer shown for the root node.

Documentation changes:
- Documentation has been added for how to extend `penzai.treescope` to render custom types in a stable way.
- The right-hand sidebar now shows more detail in tutorial notebooks.

Penzai 0.1.1 fixes a number of bugs and adds a few convenience features.

New features:
- `NamedArray` can now be converted automatically to a JAX array without unwrapping it if it has no more named axes (e.g. after calling `untag`).
- Lifting JAX functions to operate over `NamedArray`s can be accomplished with less boilerplate using `penzai.toolshed.auto_nmap`: if you set `njnp = auto_nmap.wrap_module(jnp)`, you can use `njnp.{method}` instead of `pz.nx.nmap(jnp.{method})`.

Bug fixes and improvements:
- Bugfix: Array data in array visualization tooltips now correctly accounts for slider offsets.
- Improvements to compatibility of `penzai.treescope` renderings:
- Copy-path buttons no longer require internet access.
- Pretty-printed outputs should now appear in the correct location in JupyterLab.
- Perfomance improvements for non-Chrome browsers.

Documentation changes:
- Fixed typos and formatting issues.
- Fixed notebook output display bug in Firefox.
- Source links in ReadTheDocs now link to a specific commit.
- Cross-page links in ReadTheDocs now support hover tooltip previews.

Initial release of Penzai.

Read the documentation at https://penzai.readthedocs.io/en/v0.1.0/.