Penzai 0.2.5

Bug fixes

• Fixed a broken test when running with Treescope 0.1.9 (#108)

Penzai 0.2.4

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.

Penzai 0.2.3

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.

Penzai 0.2.2

• 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)

Penzai 0.2.1

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

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

Penzai 0.1.5

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.

Penzai 0.1.4

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.

Penzai 0.1.3 (+ V2 NN API!)

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

Penzai 0.1.2

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