SYNTOLOGY HomeExplorerAtlasCodeMethodologyAboutDevelopersFeedPricing
Paper · 2408.00211 · ICML · 2024

Penzai + Treescope: A Toolkit for Interpreting, Visualizing, and Editing Models As Data

Daniel Johnson

arXiv · PDF · Open in the Atlas

Code that ran

We lifted 23 functions out of this paper's own repositories and ran 14 of them in a sandbox. "Ran" means the function executed on a synthesized input and returned a value. It is not a reproduction of the paper's results.

RepositoryRoleRan
google-deepmind/treescope canonical 10 of 17
google-deepmind/penzai canonical 4 of 6
FunctionStatusWhere it lives
color_from_string Ran google-deepmind/treescope/treescope/formatting_util.py
code served (permissive licence) · get_code("1ed46793bced4f6f")
dataclass_from_attributes Ran google-deepmind/treescope/treescope/dataclass_util.py
code served (permissive licence) · get_code("491da2a807a0d703")
get_dtype_name Ran google-deepmind/treescope/treescope/dtype_util.py
code served (permissive licence) · get_code("0bcd3e980f7fe655")
init_takes_fields Ran google-deepmind/treescope/treescope/dataclass_util.py
code served (permissive licence) · get_code("3af90f8f7873fac9")
is_floating_dtype Ran google-deepmind/treescope/treescope/dtype_util.py
code served (permissive licence) · get_code("07b1fce0fdc1eefd")
is_integer_dtype Ran google-deepmind/treescope/treescope/dtype_util.py
code served (permissive licence) · get_code("5c689b98d9651bb5")
is_pytree_dataclass_type Ran google-deepmind/penzai/penzai/core/struct.py
code served (permissive licence) · get_code("8da90bf2b11facb7")
oklch_color Ran google-deepmind/treescope/treescope/formatting_util.py
code served (permissive licence) · get_code("49862fb509edca46")
parse_simple_color_and_pattern_spec Ran google-deepmind/treescope/treescope/formatting_util.py
code served (permissive licence) · get_code("c54a7b3f3d04bf76")
styled Ran google-deepmind/treescope/treescope/figures.py
code served (permissive licence) · get_code("cc38f8434830dd01")
tree_flatten_exactly_one_level Ran google-deepmind/penzai/penzai/core/tree_util.py
code served (permissive licence) · get_code("7f0011a936783767")
var Ran google-deepmind/penzai/penzai/core/shapecheck.py
code served (permissive licence) · get_code("c716cd4fc3d93251")
vars_for_axes Ran google-deepmind/penzai/penzai/core/shapecheck.py
code served (permissive licence) · get_code("ad130894fc372aec")
with_font_size Ran google-deepmind/treescope/treescope/figures.py
code served (permissive licence) · get_code("314cb1c24954c222")
default_well_known_filter Not yet run google-deepmind/treescope/treescope/canonical_aliases.py
code served (permissive licence) · get_code("4eb59933d52ad100")
indented Not yet run google-deepmind/treescope/treescope/figures.py
code served (permissive licence) · get_code("e4fd281c5d149c61")
is_pytree_node_field Not yet run google-deepmind/penzai/penzai/core/struct.py
code served (permissive licence) · get_code("ffd28ac2c65414f3")
lookup_alias Not yet run google-deepmind/treescope/treescope/canonical_aliases.py
code served (permissive licence) · get_code("5c75d7c159ec68b6")
maybe_defer_rendering Not yet run google-deepmind/treescope/treescope/lowering.py
code served (permissive licence) · get_code("fe56c9c538b60fab")
maybe_local_module_name Not yet run google-deepmind/treescope/treescope/canonical_aliases.py
code served (permissive licence) · get_code("ad8ee3a7572721bb")
pretty_keystr Not yet run google-deepmind/penzai/penzai/core/tree_util.py
code served (permissive licence) · get_code("c45870654f09629f")
render_to_html_as_root Not yet run google-deepmind/treescope/treescope/lowering.py
code served (permissive licence) · get_code("4fe2c2ee820dc796")
render_to_text_as_root Not yet run google-deepmind/treescope/treescope/lowering.py
code served (permissive licence) · get_code("2a11ffd4903b882a")

Repositories linked to this paper

Some links come from the archived Papers with Code dataset (CC BY-SA 4.0): attribution and licence.

Abstract

Much of today's machine learning research involves interpreting, modifying or visualizing models after they are trained. I present Penzai, a neural network library designed to simplify model manipulation by representing models as simple data structures, and Treescope, an interactive prettyprinter and array visualizer that can visualize both model inputs/outputs and the models themselves. Penzai models are built using declarative combinators that expose the model forward pass in the structure of the model object itself, and use named axes to ensure each operation is semantically meaningful. With Penzai's tree-editing selector system, users can both insert and replace model components, allowing them to intervene on intermediate values or make other edits to the model structure. Users can then get immediate feedback by visualizing the modified model with Treescope. I describe the motivation and main features of Penzai and Treescope, and discuss how treating the model as data enables a variety of analyses and interventions to be implemented as data-structure transformations, without requiring model designers to add explicit hooks.

For agents

The same record, over MCP at https://syntology.ai/mcp:

get_harvested_code_for_paper("2408.00211")
get_code_for_paper("2408.00211")
have("2408.00211")

Connect an agent — have() is free.