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.