Releases: google-deepmind/treescope
Releases · google-deepmind/treescope
Treescope 0.1.6
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 theuv
package manager. - Checks for formatting (via
pyink
), linting (viapylint
), and type annotations (viapytype
) are now run on every PR. pyink
andpylint
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
Bug fixes and improvements:
- Removed check for deprecated
named_shape
attribute onjax.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
Bug fixes and improvements:
- Figure rendering and automatic visualization:
- Improved iframe size detection logic
- Added ability to collapse autovisualizer
IPythonVisualization
outputs whenreplace=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
New features
render_array
andArrayAutovisualizer
now accept apixels_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
inArrayAutovisualizer
. - 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
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
New features:
- Added the
autovisualize
keyword argument totreescope.show
andtreescope.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
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 intotreescope
. - 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).