Skip to content

Releases: google-deepmind/treescope

Treescope 0.1.6

21 Nov 23:06
Compare
Choose a tag to compare

New features

  • Rendering JAX arrays is now much faster, especially for small device arrays (#40)

Bug fixes and improvements

  • Fixed incorrect bloat16 array formatting (#43)
  • Fixed a number of typos (#44)
  • Fixed render_array to correctly support re-ordering rows and columns to match documented behavior (#49)
  • Test no longer depend on the exact representation of jax.tree_util.GetAttrKey, making it robust to future planned changes in JAX (#48)

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.

Treescope 0.1.5

13 Sep 20:38
Compare
Choose a tag to compare

Bug fixes and improvements:

  • Removed check for deprecated named_shape attribute on jax.ShapeDtypeStruct, since this attribute is being removed from JAX.
  • Added a "part" attribute to the root rendering element to allow customizing CSS styles via the CSS shadow parts system (#38).

Treescope 0.1.4

06 Sep 02:18
Compare
Choose a tag to compare

Bug fixes and improvements:

  • Figure rendering and automatic visualization:
    • Improved iframe size detection logic
    • Added ability to collapse autovisualizer IPythonVisualization outputs when replace=False
    • Reduced iframe duplication when using Treescope figures inside autovisualizers
    • Ensured DOM nodes are preserved when substituting deferred renderings to avoid breaking event handlers in figures
  • Added support for rendering JAX arrays with PmapSharding shardings
  • Improved support for rendering PyTorch tensors:
    • Turned off summarization for tensors on the "meta" device, since they have no data
    • Added support for float16 and bfloat16 tensors
    • Added device annotations in tensor summaries
  • Added default support for collections.UserDict and subclasses

Treescope 0.1.3

27 Aug 21:57
Compare
Choose a tag to compare

New features

  • render_array and ArrayAutovisualizer now accept a pixels_per_cell argument to customize the initial zoom level of array renderings.
  • Array renderings now adapt their size and rendering options based on device resolution, to ensure renderings are crisp when the window is zoomed in or when used on a high-DPI display.

Bug fixes and improvements:

  • Fixed a type error with prefers_row/prefers_column in ArrayAutovisualizer.
  • Simplified torch array adapter to not require named axis support or isposinf/isneginf primitives, allowing it to be used with the "mps" backend (Apple silicon).
  • Treescope now correctly catches and handles errors that occur during deferred rendering or inside custom __repr__ implementations.

Treescope 0.1.2

12 Aug 20:13
Compare
Choose a tag to compare

Bug fixes and improvements:

  • Fixed rendering of PyTorch GPU tensors (#19)
  • Treescope now avoids overriding the custom IPython format function _repr_mimebundle_ for types that define it

Documentation changes:

  • Fixed typos in README and docs home page (#22)
  • Added information about how to cite the accompanying whitepaper

Treescope 0.1.1

07 Aug 01:04
Compare
Choose a tag to compare

New features:

  • Added the autovisualize keyword argument to treescope.show and treescope.display to allow customizing the automatic visualizer on a per-call basis, along with a few more configuration options.
  • Added alternate forms %%autovisualize True and %%autovisualize False for the %%autovisualize magic.

Bug fixes and improvements:

  • Adjusted output wrapper to ensure that "Loading" indicators appear when the browser is loading a rendered output.
  • Fix a visual bug with torch.nn.Sequential rendering.

Documentation changes:

  • Added documentation of neural network support.
  • Fixed broken links.

Treescope 0.1.0

25 Jul 22:48
Compare
Choose a tag to compare

Initial release of the standalone Treescope pretty-printer.

Treescope was originally designed as the pretty-printer for the Penzai neural network library, but has been moved to its own repository so that it can be installed and used independently. See the commit history for Penzai for earlier versions of Treescope.

Changes relative to the penzai.treescope submodule of Penzai release 0.1.5:

  • Moved all functions and methods out of the penzai.treescope namespace and into treescope.
  • Renamed extension protocol method from __penzai_repr__ to __treescope_repr__.
  • Added a global handler registry, to enable registering handlers for existing types.
  • Redesigned the array visualization system in terms of an NDArrayAdapter protocol, enabling it to support non-JAX arrays.
  • Added handlers for PyTorch tensors and PyTorch neural network modules.
  • Moved implementation details of the internal intermediate representation into an _internal subpackage, and added new public construction methods for various rendering components.
  • Removed some complex features that were rarely used (copy-paste fallbacks and between-node hyperlinks).