{"about":{"site":"https://codewithpapers.app","non_affiliation":"Code with Papers and Syntology are not affiliated with, endorsed by, or sponsored by Papers with Code, Meta, or the pwc-archive mirror.","licence":"CC BY-SA 4.0","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","attribution":"https://codewithpapers.app/attribution","modified":"archive material modified by Syntology; see the attribution page"},"url":"/paper/penzai-treescope-a-toolkit-for-interpreting","title":"Penzai + Treescope: A Toolkit for Interpreting, Visualizing, and Editing Models As Data","arxiv_id":"2408.00211","date":"2024-08-01","proceeding":null,"authors":["Daniel D. Johnson"],"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 pretty-printer 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.","url_abs":"https://arxiv.org/abs/2408.00211v1","url_pdf":"https://arxiv.org/pdf/2408.00211v1.pdf","source":{"archive":"pwc-archive (Hugging Face), CC BY-SA 4.0","snapshot":"2025-07-28","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","row_kind":"abstracts"},"code_links":[{"paper_slug":"penzai-treescope-a-toolkit-for-interpreting","repo_url":"https://github.com/google-deepmind/penzai","is_official":1,"mentioned_in_paper":0,"mentioned_in_github":1,"framework":"jax","reach":{"status":"ok","spdx":"Apache-2.0"}},{"paper_slug":"penzai-treescope-a-toolkit-for-interpreting","repo_url":"https://github.com/google-deepmind/treescope","is_official":1,"mentioned_in_paper":0,"mentioned_in_github":1,"framework":"pytorch","reach":{"status":"ok","spdx":"Apache-2.0"}}],"tasks":[],"methods":[{"method_slug":null,"method_name":"Library"}],"datasets_introduced":[],"methods_introduced":[],"results":[],"syntology":{"syntology_url":"https://syntology.ai/paper/2408.00211","atlas_url":"https://app.syntology.ai/?focus=2408.00211","mcp":{"get_harvested_code_for_paper":{"arxiv_id":"2408.00211"}},"developers":"https://syntology.ai/developers","read_at":"2026-09-25T09:33:49+00:00","read_at_is":"when the build read Syntology's graph, not when any sample ran","claim":"Per-sample execution status on synthesized fixtures; not a correctness claim about the paper. Samples come from repositories linked to the paper, official or community; repo_kind says which.","repos":[{"provenance":"deterministic:regex_extraction","url":"https://github.com/google-deepmind/penzai","reach":{"status":"ok","spdx":"Apache-2.0"}},{"provenance":"deterministic:regex_extraction","url":"https://github.com/google-deepmind/treescope","reach":{"status":"ok","spdx":"Apache-2.0"}}],"summary":{"ran":14,"unverified":9},"by_repo_kind":{"official":{"samples":23,"ran":14,"repositories":2}},"repo_kind_vocabulary":{"official":"The archive marks this repository official for the paper","named_in_paper":"The archive records that the paper mentions this repository; it is not marked official","listed":"In the archive's code links for this paper, not marked official and not recorded as mentioned in the paper","found_in_text":"Syntology found this repository in the paper's own text; whether it is the authors' implementation is not asserted","community":"Not in the archive's code links for this paper; a community repository Syntology harvested"},"n_pointer_only_for_licence":0,"samples":[{"code_sha256_prefix":"1ed46793bced4f6f","entry":"color_from_string","repo":"google-deepmind/treescope","repo_kind":"official","path":"treescope/formatting_util.py","file_url":"https://github.com/google-deepmind/treescope/blob/HEAD/treescope/formatting_util.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"1ed46793bced4f6f"}},{"code_sha256_prefix":"491da2a807a0d703","entry":"dataclass_from_attributes","repo":"google-deepmind/treescope","repo_kind":"official","path":"treescope/dataclass_util.py","file_url":"https://github.com/google-deepmind/treescope/blob/HEAD/treescope/dataclass_util.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"491da2a807a0d703"}},{"code_sha256_prefix":"0bcd3e980f7fe655","entry":"get_dtype_name","repo":"google-deepmind/treescope","repo_kind":"official","path":"treescope/dtype_util.py","file_url":"https://github.com/google-deepmind/treescope/blob/HEAD/treescope/dtype_util.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"0bcd3e980f7fe655"}},{"code_sha256_prefix":"3af90f8f7873fac9","entry":"init_takes_fields","repo":"google-deepmind/treescope","repo_kind":"official","path":"treescope/dataclass_util.py","file_url":"https://github.com/google-deepmind/treescope/blob/HEAD/treescope/dataclass_util.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"3af90f8f7873fac9"}},{"code_sha256_prefix":"07b1fce0fdc1eefd","entry":"is_floating_dtype","repo":"google-deepmind/treescope","repo_kind":"official","path":"treescope/dtype_util.py","file_url":"https://github.com/google-deepmind/treescope/blob/HEAD/treescope/dtype_util.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"07b1fce0fdc1eefd"}},{"code_sha256_prefix":"5c689b98d9651bb5","entry":"is_integer_dtype","repo":"google-deepmind/treescope","repo_kind":"official","path":"treescope/dtype_util.py","file_url":"https://github.com/google-deepmind/treescope/blob/HEAD/treescope/dtype_util.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"5c689b98d9651bb5"}},{"code_sha256_prefix":"8da90bf2b11facb7","entry":"is_pytree_dataclass_type","repo":"google-deepmind/penzai","repo_kind":"official","path":"penzai/core/struct.py","file_url":"https://github.com/google-deepmind/penzai/blob/HEAD/penzai/core/struct.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"8da90bf2b11facb7"}},{"code_sha256_prefix":"49862fb509edca46","entry":"oklch_color","repo":"google-deepmind/treescope","repo_kind":"official","path":"treescope/formatting_util.py","file_url":"https://github.com/google-deepmind/treescope/blob/HEAD/treescope/formatting_util.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"49862fb509edca46"}},{"code_sha256_prefix":"c54a7b3f3d04bf76","entry":"parse_simple_color_and_pattern_spec","repo":"google-deepmind/treescope","repo_kind":"official","path":"treescope/formatting_util.py","file_url":"https://github.com/google-deepmind/treescope/blob/HEAD/treescope/formatting_util.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"c54a7b3f3d04bf76"}},{"code_sha256_prefix":"cc38f8434830dd01","entry":"styled","repo":"google-deepmind/treescope","repo_kind":"official","path":"treescope/figures.py","file_url":"https://github.com/google-deepmind/treescope/blob/HEAD/treescope/figures.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"cc38f8434830dd01"}},{"code_sha256_prefix":"7f0011a936783767","entry":"tree_flatten_exactly_one_level","repo":"google-deepmind/penzai","repo_kind":"official","path":"penzai/core/tree_util.py","file_url":"https://github.com/google-deepmind/penzai/blob/HEAD/penzai/core/tree_util.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"7f0011a936783767"}},{"code_sha256_prefix":"c716cd4fc3d93251","entry":"var","repo":"google-deepmind/penzai","repo_kind":"official","path":"penzai/core/shapecheck.py","file_url":"https://github.com/google-deepmind/penzai/blob/HEAD/penzai/core/shapecheck.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"c716cd4fc3d93251"}},{"code_sha256_prefix":"ad130894fc372aec","entry":"vars_for_axes","repo":"google-deepmind/penzai","repo_kind":"official","path":"penzai/core/shapecheck.py","file_url":"https://github.com/google-deepmind/penzai/blob/HEAD/penzai/core/shapecheck.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"ad130894fc372aec"}},{"code_sha256_prefix":"314cb1c24954c222","entry":"with_font_size","repo":"google-deepmind/treescope","repo_kind":"official","path":"treescope/figures.py","file_url":"https://github.com/google-deepmind/treescope/blob/HEAD/treescope/figures.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"314cb1c24954c222"}},{"code_sha256_prefix":"4eb59933d52ad100","entry":"default_well_known_filter","repo":"google-deepmind/treescope","repo_kind":"official","path":"treescope/canonical_aliases.py","file_url":"https://github.com/google-deepmind/treescope/blob/HEAD/treescope/canonical_aliases.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"4eb59933d52ad100"}},{"code_sha256_prefix":"e4fd281c5d149c61","entry":"indented","repo":"google-deepmind/treescope","repo_kind":"official","path":"treescope/figures.py","file_url":"https://github.com/google-deepmind/treescope/blob/HEAD/treescope/figures.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"e4fd281c5d149c61"}},{"code_sha256_prefix":"ffd28ac2c65414f3","entry":"is_pytree_node_field","repo":"google-deepmind/penzai","repo_kind":"official","path":"penzai/core/struct.py","file_url":"https://github.com/google-deepmind/penzai/blob/HEAD/penzai/core/struct.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"ffd28ac2c65414f3"}},{"code_sha256_prefix":"5c75d7c159ec68b6","entry":"lookup_alias","repo":"google-deepmind/treescope","repo_kind":"official","path":"treescope/canonical_aliases.py","file_url":"https://github.com/google-deepmind/treescope/blob/HEAD/treescope/canonical_aliases.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"5c75d7c159ec68b6"}},{"code_sha256_prefix":"fe56c9c538b60fab","entry":"maybe_defer_rendering","repo":"google-deepmind/treescope","repo_kind":"official","path":"treescope/lowering.py","file_url":"https://github.com/google-deepmind/treescope/blob/HEAD/treescope/lowering.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"fe56c9c538b60fab"}},{"code_sha256_prefix":"ad8ee3a7572721bb","entry":"maybe_local_module_name","repo":"google-deepmind/treescope","repo_kind":"official","path":"treescope/canonical_aliases.py","file_url":"https://github.com/google-deepmind/treescope/blob/HEAD/treescope/canonical_aliases.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"ad8ee3a7572721bb"}},{"code_sha256_prefix":"c45870654f09629f","entry":"pretty_keystr","repo":"google-deepmind/penzai","repo_kind":"official","path":"penzai/core/tree_util.py","file_url":"https://github.com/google-deepmind/penzai/blob/HEAD/penzai/core/tree_util.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"c45870654f09629f"}},{"code_sha256_prefix":"4fe2c2ee820dc796","entry":"render_to_html_as_root","repo":"google-deepmind/treescope","repo_kind":"official","path":"treescope/lowering.py","file_url":"https://github.com/google-deepmind/treescope/blob/HEAD/treescope/lowering.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"4fe2c2ee820dc796"}},{"code_sha256_prefix":"2a11ffd4903b882a","entry":"render_to_text_as_root","repo":"google-deepmind/treescope","repo_kind":"official","path":"treescope/lowering.py","file_url":"https://github.com/google-deepmind/treescope/blob/HEAD/treescope/lowering.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"Apache-2.0","inline_ok":true,"mcp_get_code":{"code_sha256":"2a11ffd4903b882a"}}]},"arxiv_metadata":null,"syntology_extracted_results":null}