Skip to content

Penzai 0.2.0
Compare
Choose a tag to compare
@danieldjohnson danieldjohnson released this 26 Jul 23:39
· 17 commits to main since this release
v0.2.0
a4fbb59

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.
Assets 2
Loading