├── penzai ├── nn │ ├── __init__.py │ ├── _treescope_handlers │ │ ├── __init__.py │ │ └── layer_handler.py │ ├── basic_ops.py │ └── combinators.py ├── models │ ├── __init__.py │ ├── transformer │ │ ├── __init__.py │ │ ├── variants │ │ │ ├── __init__.py │ │ │ ├── llama.py │ │ │ └── mistral.py │ │ └── simple_decoding_loop.py │ └── simple_mlp.py ├── toolshed │ ├── __init__.py │ ├── jit_wrapper.py │ ├── patch_ipdb.py │ ├── gradient_checkpointing.py │ └── auto_nmap.py ├── deprecated │ ├── v1 │ │ ├── core │ │ │ ├── _treescope_handlers │ │ │ │ └── __init__.py │ │ │ ├── __init__.py │ │ │ └── random_stream.py │ │ ├── __init__.py │ │ ├── example_models │ │ │ ├── __init__.py │ │ │ ├── gemma │ │ │ │ └── __init__.py │ │ │ └── simple_mlp.py │ │ ├── data_effects │ │ │ └── __init__.py │ │ ├── nn │ │ │ ├── __init__.py │ │ │ ├── basic_ops.py │ │ │ └── combinators.py │ │ ├── pz │ │ │ ├── de.py │ │ │ ├── nn.py │ │ │ └── __init__.py │ │ └── toolshed │ │ │ ├── jit_wrapper.py │ │ │ └── sharding_util.py │ └── __init__.py ├── experimental │ └── __init__.py ├── core │ ├── _treescope_handlers │ │ └── __init__.py │ ├── __init__.py │ ├── syntactic_sugar.py │ ├── random_stream.py │ └── tree_util.py ├── __init__.py ├── treescope │ ├── formatting_util.py │ ├── repr_lib.py │ └── __init__.py └── pz │ ├── ts.py │ ├── nn.py │ └── __init__.py ├── docs ├── .gitignore ├── _static │ ├── readme_teaser.png │ ├── custom.css │ └── custom.js ├── _templates │ ├── pzbase.rst │ ├── pzdata.rst │ ├── pzmodule.rst │ ├── pzmodule_full.rst │ └── pzclass.rst ├── Makefile ├── _autogen_root.rst ├── api │ ├── treescope.rst │ ├── pz.ts.rst │ ├── pz.nn.rst │ └── pz.rst ├── ext │ ├── pz_alias_rewrite.py │ └── nb_output_cell_to_iframe.py └── scripts │ └── readthedocs_fetch_notebook_outputs.sh ├── .vscode ├── extensions.json └── settings.json ├── AUTHORS ├── .gitignore ├── tests ├── __init__.py ├── models │ ├── __init__.py │ └── simple_mlp_test.py ├── nn │ ├── __init__.py │ ├── basic_ops_test.py │ ├── embedding_test.py │ ├── standardization_test.py │ ├── parameters_test.py │ ├── layer_test.py │ └── grouping_test.py ├── deprecated │ ├── __init__.py │ └── v1 │ │ ├── __init__.py │ │ ├── nn │ │ ├── __init__.py │ │ ├── basic_ops_test.py │ │ ├── embedding_test.py │ │ ├── standardization_test.py │ │ ├── grouping_test.py │ │ └── linear_and_affine_test.py │ │ ├── data_effects │ │ ├── __init__.py │ │ ├── side_output_test.py │ │ └── random_test.py │ │ ├── toolshed │ │ ├── __init__.py │ │ └── lora_test.py │ │ ├── example_models │ │ ├── __init__.py │ │ └── simple_mlp_test.py │ │ ├── misc_util_test.py │ │ └── shapecheck_layer_test.py ├── toolshed │ ├── __init__.py │ ├── jit_wrapper_test.py │ ├── gradient_checkpointing_test.py │ ├── auto_nmap_test.py │ ├── save_intermediates_test.py │ └── lora_test.py ├── treescope │ └── fixtures │ │ ├── __init__.py │ │ └── treescope_examples_fixture.py └── core │ ├── auto_order_types_test.py │ └── misc_util_test.py ├── .readthedocs.yaml ├── run_tests.py ├── CONTRIBUTING.md ├── .github └── workflows │ └── unittests.yml └── pyproject.toml /penzai/nn/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /penzai/models/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /penzai/toolshed/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /penzai/nn/_treescope_handlers/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /docs/.gitignore: -------------------------------------------------------------------------------- 1 | **/_autosummary 2 | _build 3 | -------------------------------------------------------------------------------- /penzai/deprecated/v1/core/_treescope_handlers/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /docs/_static/readme_teaser.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/penzai/HEAD/docs/_static/readme_teaser.png -------------------------------------------------------------------------------- /.vscode/extensions.json: -------------------------------------------------------------------------------- 1 | { 2 | "recommendations": [ 3 | "ms-python.black-formatter", 4 | "ms-python.pylint" 5 | ] 6 | } 7 | -------------------------------------------------------------------------------- /docs/_templates/pzbase.rst: -------------------------------------------------------------------------------- 1 | {{ name | escape | underline}} 2 | 3 | .. currentmodule:: {{ module }} 4 | 5 | .. auto{{ objtype }}:: {{ objname }} 6 | -------------------------------------------------------------------------------- /docs/_templates/pzdata.rst: -------------------------------------------------------------------------------- 1 | {{ name | escape | underline}} 2 | 3 | .. currentmodule:: {{ module }} 4 | 5 | .. autodata:: {{ objname }} 6 | :no-value: 7 | -------------------------------------------------------------------------------- /AUTHORS: -------------------------------------------------------------------------------- 1 | # This is the list of Penzai's significant contributors. 2 | # 3 | # This does not necessarily list everyone who has contributed code, 4 | # especially since many employees of one corporation may be contributing. 5 | # To see the full list of contributors, see the revision history in 6 | # source control. 7 | DeepMind Technologies Limited 8 | Daniel D. Johnson 9 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Compiled python modules. 2 | *.pyc 3 | 4 | # Byte-compiled 5 | _pycache__/ 6 | .cache/ 7 | 8 | # Poetry, setuptools, PyPI distribution artifacts. 9 | /*.egg-info 10 | .eggs/ 11 | build/ 12 | dist/ 13 | poetry.lock 14 | 15 | # Tests 16 | .pytest_cache/ 17 | 18 | # Type checking 19 | .pytype/ 20 | 21 | # Virtual env 22 | .venv/ 23 | 24 | # Other 25 | *.DS_Store 26 | 27 | # PyCharm 28 | .idea 29 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /tests/models/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /tests/nn/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /penzai/deprecated/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /tests/deprecated/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /tests/toolshed/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /penzai/deprecated/v1/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /penzai/experimental/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /tests/deprecated/v1/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /tests/deprecated/v1/nn/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /penzai/deprecated/v1/core/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /tests/treescope/fixtures/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /tests/deprecated/v1/data_effects/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /tests/deprecated/v1/toolshed/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /penzai/deprecated/v1/example_models/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /tests/deprecated/v1/example_models/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /penzai/core/_treescope_handlers/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | """Treescope handlers for Penzai's own classes.""" 16 | -------------------------------------------------------------------------------- /docs/Makefile: -------------------------------------------------------------------------------- 1 | # Minimal makefile for Sphinx documentation 2 | # 3 | # You can set these variables from the command line. 4 | SPHINXOPTS = 5 | SPHINXBUILD = sphinx-build 6 | SOURCEDIR = . 7 | BUILDDIR = _build 8 | 9 | # Put it first so that "make" without argument is like "make help". 10 | help: 11 | @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 12 | 13 | .PHONY: help Makefile 14 | 15 | # Catch-all target: route all unknown targets to Sphinx using the new 16 | # "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). 17 | %: Makefile 18 | @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 19 | 20 | clean: 21 | rm -rf $(BUILDDIR)/* _collections _autosummary modules 22 | -------------------------------------------------------------------------------- /penzai/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | """A JAX research toolkit for building, editing and visualizing neural networks.""" # pylint: disable=line-too-long 16 | 17 | __version__ = '0.2.5' 18 | -------------------------------------------------------------------------------- /.readthedocs.yaml: -------------------------------------------------------------------------------- 1 | # Read the Docs configuration file 2 | # See https://docs.readthedocs.io/en/stable/config-file/v2.html for details 3 | 4 | version: 2 5 | 6 | build: 7 | os: ubuntu-22.04 8 | tools: 9 | python: "3.11" 10 | commands: 11 | # Fetch precomputed notebook outputs (if any) 12 | - bash docs/scripts/readthedocs_fetch_notebook_outputs.sh 13 | # Install and build using uv 14 | - asdf plugin add uv 15 | - asdf install uv latest 16 | - asdf global uv latest 17 | - uv sync --extra docs --frozen 18 | - uv pip install readthedocs-sphinx-ext 19 | - cd docs && uv run python -m sphinx -T -b html -d docs/_build/doctrees -D language=en . $READTHEDOCS_OUTPUT/html 20 | 21 | sphinx: 22 | builder: html 23 | configuration: docs/conf.py 24 | fail_on_warning: false 25 | -------------------------------------------------------------------------------- /penzai/models/transformer/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | """Basic transformer backbone with variants.""" 16 | 17 | from . import model_parts 18 | from . import sampling_mode 19 | from . import simple_decoding_loop 20 | from . import variants 21 | -------------------------------------------------------------------------------- /penzai/deprecated/v1/example_models/gemma/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | """Modular implementation of the Gemma model architecture.""" 16 | 17 | from . import model_core 18 | from . import sampling_mode 19 | from . import simple_decoding_loop 20 | -------------------------------------------------------------------------------- /penzai/models/transformer/variants/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | """Concrete transformer variants.""" 16 | 17 | from . import gemma 18 | from . import gpt_neox 19 | from . import llama 20 | from . import llamalike_common 21 | from . import mistral 22 | -------------------------------------------------------------------------------- /penzai/deprecated/v1/data_effects/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | """Penzai's data effect system.""" 16 | 17 | from . import effect_base 18 | from . import local_state 19 | from . import random 20 | from . import side_input 21 | from . import side_output 22 | -------------------------------------------------------------------------------- /docs/_autogen_root.rst: -------------------------------------------------------------------------------- 1 | .. 2 | This file is not actually referenced in the docs, but it is the entry point 3 | that generates automatic summaries for Penzai's modules. `index.rst` points 4 | directly at the autosummary files generated while processing this one. 5 | We also reference other orphaned modules here so that Sphinx doesn't give 6 | warnings about them. 7 | 8 | :orphan: 9 | 10 | .. autosummary:: 11 | :toctree: _autosummary 12 | :template: pzmodule_full.rst 13 | :recursive: 14 | 15 | penzai.core 16 | penzai.nn 17 | penzai.models 18 | penzai.toolshed 19 | 20 | penzai.deprecated.v1.core 21 | penzai.deprecated.v1.nn 22 | penzai.deprecated.v1.data_effects 23 | penzai.deprecated.v1.example_models 24 | penzai.deprecated.v1.toolshed 25 | 26 | .. toctree:: 27 | :hidden: 28 | 29 | notebooks/induction_heads_2B 30 | _include/_glue_figures 31 | -------------------------------------------------------------------------------- /penzai/core/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | """Penzai core types and functions.""" 16 | 17 | from . import named_axes 18 | from . import partitioning 19 | from . import selectors 20 | from . import shapecheck 21 | from . import struct 22 | from . import syntactic_sugar 23 | from . import tree_util 24 | -------------------------------------------------------------------------------- /run_tests.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """Entry point executable to run all tests.""" 15 | 16 | import subprocess 17 | 18 | if __name__ == "__main__": 19 | subprocess.check_call( 20 | ["python", "-m", "pytest", "tests", "-k", "not ShardingUtilTest"] 21 | ) 22 | subprocess.check_call( 23 | ["python", "-m", "pytest", "tests", "-k", "ShardingUtilTest"] 24 | ) 25 | -------------------------------------------------------------------------------- /penzai/deprecated/v1/nn/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | """Penzai's declarative neural network library.""" 16 | 17 | from . import attention 18 | from . import basic_ops 19 | from . import combinators 20 | from . import dropout 21 | from . import embeddings 22 | from . import grouping 23 | from . import linear_and_affine 24 | from . import parameters 25 | from . import standardization 26 | -------------------------------------------------------------------------------- /.vscode/settings.json: -------------------------------------------------------------------------------- 1 | { 2 | "files.insertFinalNewline": true, 3 | "files.trimFinalNewlines": true, 4 | "files.trimTrailingWhitespace": true, 5 | "files.associations": { 6 | ".pylintrc": "ini" 7 | }, 8 | "files.watcherExclude": { 9 | "**/.git/**": true 10 | }, 11 | "files.exclude": { 12 | "**/__pycache__": true, 13 | "**/.pytest_cache": true, 14 | "**/*.egg-info": true 15 | }, 16 | "python.testing.unittestEnabled": false, 17 | "python.testing.nosetestsEnabled": false, 18 | "python.testing.pytestEnabled": true, 19 | "python.linting.pylintUseMinimalCheckers": false, 20 | "[python]": { 21 | "editor.rulers": [80], 22 | "editor.tabSize": 2, 23 | "editor.formatOnSave": true, 24 | "editor.detectIndentation": false, 25 | "editor.defaultFormatter": "ms-python.black-formatter", 26 | }, 27 | "black-formatter.path": ["uvx", "--python=>=3.9,!=3.12.5", "pyink"], 28 | "pylint.enabled": true, 29 | } 30 | -------------------------------------------------------------------------------- /penzai/treescope/formatting_util.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | """Compatibility stub: Utilities for formatting and rendering. 16 | 17 | New code should instead use `treescope.formatting_util`. This module will be 18 | deprecated and eventually removed. 19 | """ 20 | 21 | # pylint: disable=g-importing-member,g-multiple-import,unused-import 22 | 23 | from treescope.formatting_util import ( 24 | color_from_string, 25 | oklch_color, 26 | ) 27 | -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # How to Contribute 2 | 3 | ## Contributor License Agreement 4 | 5 | Contributions to this project must be accompanied by a Contributor License 6 | Agreement. You (or your employer) retain the copyright to your contribution, 7 | this simply gives us permission to use and redistribute your contributions as 8 | part of the project. Head over to to see 9 | your current agreements on file or to sign a new one. 10 | 11 | You generally only need to submit a CLA once, so if you've already submitted one 12 | (even if it was for a different project), you probably don't need to do it 13 | again. 14 | 15 | ## Code reviews 16 | 17 | All submissions, including submissions by project members, require review. We 18 | use GitHub pull requests for this purpose. Consult 19 | [GitHub Help](https://help.github.com/articles/about-pull-requests/) for more 20 | information on using pull requests. 21 | 22 | ## Community Guidelines 23 | 24 | This project follows [Google's Open Source Community 25 | Guidelines](https://opensource.google/conduct/). 26 | -------------------------------------------------------------------------------- /penzai/treescope/repr_lib.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | """Compatibility stub: Stable high-level interface for building object reprs. 16 | 17 | New code should instead use `treescope.repr_lib`. This module will be 18 | deprecated and eventually removed. 19 | """ 20 | 21 | # pylint: disable=g-importing-member,g-multiple-import,unused-import 22 | 23 | from treescope.repr_lib import ( 24 | render_object_constructor, 25 | render_dictionary_wrapper, 26 | render_enumlike_item, 27 | ) 28 | -------------------------------------------------------------------------------- /penzai/treescope/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | """Compatibility stub for treescope. 16 | 17 | The treescope pretty-printer has moved into a separate module `treescope`. 18 | This module contains compatibility stubs for the previously-stable parts of the 19 | extension API, which are still available in the `penzai.treescope` namespace. 20 | 21 | New code should import from `treescope` directly. This module and its submodules 22 | will be deprecated and eventually removed. 23 | """ 24 | 25 | from . import formatting_util 26 | from . import repr_lib 27 | -------------------------------------------------------------------------------- /docs/api/treescope.rst: -------------------------------------------------------------------------------- 1 | penzai.treescope (Moved!) 2 | ========================================================= 3 | 4 | .. module:: penzai.treescope 5 | 6 | Treescope, Penzai's interactive pretty-printer, has moved to an independent 7 | package `treescope`. See the 8 | `Treescope documentation `_ 9 | for more information. 10 | 11 | The subpackage `penzai.treescope` contains compatibility stubs to ensure that 12 | libraries that use `penzai.treescope` to render custom types continue to work. 13 | We recommend that new users use `treescope` directly. 14 | 15 | One exception is if you want to use Treescope to render custom types that were 16 | configured using the old ``__penzai_repr__`` extension method and which have 17 | not yet been migrated to use the ``__treescope_repr__`` method. In this case, 18 | you can enable Treescope using the legacy Penzai API, using :: 19 | 20 | pz.ts.basic_interactive_setup() 21 | 22 | or, for more control: :: 23 | 24 | pz.ts.register_as_default() 25 | pz.ts.register_autovisualize_magic() 26 | pz.ts.active_autovisualizer.set_globally(pz.ts.ArrayAutovisualizer()) 27 | 28 | You can also pretty-print individual values using `pz.show` or `pz.ts.display`. 29 | -------------------------------------------------------------------------------- /tests/deprecated/v1/nn/basic_ops_test.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | """Tests for basic ops.""" 16 | 17 | from absl.testing import absltest 18 | import jax 19 | from penzai.deprecated.v1 import pz 20 | from penzai.deprecated.v1.toolshed import check_layers_by_tracing 21 | 22 | 23 | class BasicOpsTest(absltest.TestCase): 24 | 25 | def test_elementwise(self): 26 | layer = pz.nn.Elementwise(jax.nn.relu) 27 | check_layers_by_tracing.check_layer(layer) 28 | 29 | def test_softmax(self): 30 | layer = pz.nn.Softmax(("foo", "bar")) 31 | check_layers_by_tracing.check_layer(layer) 32 | 33 | 34 | if __name__ == "__main__": 35 | absltest.main() 36 | -------------------------------------------------------------------------------- /penzai/core/syntactic_sugar.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | """Syntactic sugar utilities for common operations.""" 16 | 17 | import dataclasses 18 | 19 | 20 | @dataclasses.dataclass(frozen=True) 21 | class SliceLike: 22 | """Builds a slice when sliced (e.g. ``pz.slice[1:3] == slice(1, 3, None)``). 23 | 24 | An instance of this class is exposed as `pz.slice`, so that you can easily 25 | slice named array axes using syntax like ``arr[{"name": pz.slice[1:3]}]`` 26 | instead of ``arr[{"name": slice(1, 3, None)}]``. 27 | """ 28 | 29 | def __getitem__(self, indexer): 30 | return indexer 31 | 32 | 33 | slice = SliceLike() # pylint: disable=redefined-builtin 34 | -------------------------------------------------------------------------------- /docs/_static/custom.css: -------------------------------------------------------------------------------- 1 | .cell_output .output.text_html, .glued-cell-output .output.text_html { 2 | background-color: white; 3 | line-height: 1.24; 4 | } 5 | 6 | iframe.cell_output_frame { 7 | width: 100%; 8 | border: 1px solid #7f7f7f2e; 9 | width: 100%; 10 | min-width: 100%; 11 | max-height: 50em; 12 | overflow: scroll; 13 | background-color: white; 14 | position: relative; 15 | } 16 | 17 | html { 18 | scroll-behavior: auto !important; 19 | } 20 | 21 | .title.logo__title { 22 | font-family:monospace; 23 | } 24 | .title.logo__title::before { 25 | color: #cccccc; 26 | position: relative; 27 | left: -0.75ch; 28 | content: '▼'; 29 | } 30 | 31 | html[data-theme=light], html[data-theme=dark] { 32 | --pst-color-inline-code: var(--pst-color-text-base) !important; 33 | } 34 | 35 | @media (min-width: 992px) { 36 | .bd-page-width { 37 | max-width: 100% !important; 38 | } 39 | .bd-sidebar-primary { 40 | flex-basis: calc(min(20%, 25em)) !important; 41 | } 42 | } 43 | 44 | html[data-theme="light"] .highlight .gd { 45 | color: #A00000 !important; 46 | } 47 | html[data-theme="light"] .highlight .gi { 48 | color: #008400 !important; 49 | } 50 | html[data-theme="dark"] .highlight .gd { 51 | color: #fe7b6e !important; 52 | } 53 | html[data-theme="dark"] .highlight .gi { 54 | color: #56d364 !important; 55 | } 56 | -------------------------------------------------------------------------------- /tests/nn/basic_ops_test.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | """Tests for basic ops.""" 16 | 17 | from absl.testing import absltest 18 | import jax 19 | from penzai import pz 20 | 21 | 22 | class BasicOpsTest(absltest.TestCase): 23 | 24 | def test_elementwise(self): 25 | layer = pz.nn.Elementwise(jax.nn.relu) 26 | result = layer(pz.nx.ones({"foo": 3})) 27 | pz.chk.check_structure(result, pz.chk.ArraySpec(named_shape={"foo": 3})) 28 | 29 | def test_softmax(self): 30 | layer = pz.nn.Softmax(("foo", "bar")) 31 | result = layer(pz.nx.ones({"foo": 3, "bar": 4, "baz": 5})) 32 | pz.chk.check_structure( 33 | result, pz.chk.ArraySpec(named_shape={"foo": 3, "bar": 4, "baz": 5}) 34 | ) 35 | 36 | 37 | if __name__ == "__main__": 38 | absltest.main() 39 | -------------------------------------------------------------------------------- /tests/deprecated/v1/nn/embedding_test.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | """Tests for embedding layer.""" 16 | 17 | from absl.testing import absltest 18 | from penzai.deprecated.v1 import pz 19 | from penzai.deprecated.v1.toolshed import check_layers_by_tracing 20 | 21 | 22 | class EmbeddingTest(absltest.TestCase): 23 | 24 | def test_embedding_lookup(self): 25 | layer = pz.nn.EmbeddingLookup( 26 | pz.nn.EmbeddingTable.from_config( 27 | vocab_size=17, embedding_axes={"foo": 5, "bar": 7} 28 | ) 29 | ) 30 | check_layers_by_tracing.check_layer(layer) 31 | 32 | def test_embedding_decode(self): 33 | layer = pz.nn.EmbeddingDecode( 34 | pz.nn.EmbeddingTable.from_config( 35 | vocab_size=17, embedding_axes={"foo": 5, "bar": 7} 36 | ) 37 | ) 38 | check_layers_by_tracing.check_layer(layer) 39 | 40 | 41 | if __name__ == "__main__": 42 | absltest.main() 43 | -------------------------------------------------------------------------------- /tests/deprecated/v1/nn/standardization_test.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | """Tests for standardization layers.""" 16 | 17 | from absl.testing import absltest 18 | from penzai.deprecated.v1 import pz 19 | from penzai.deprecated.v1.toolshed import check_layers_by_tracing 20 | 21 | 22 | class StandardizationTest(absltest.TestCase): 23 | 24 | def test_standardize_one(self): 25 | layer = pz.nn.Standardize(across="foo") 26 | check_layers_by_tracing.check_layer(layer) 27 | 28 | def test_standardize_multi(self): 29 | layer = pz.nn.Standardize(across=("foo", "bar")) 30 | check_layers_by_tracing.check_layer(layer) 31 | 32 | def test_layernorm(self): 33 | layer = pz.nn.LayerNorm.from_config(across_axes={"foo": 5, "bar": 7}) 34 | check_layers_by_tracing.check_layer( 35 | layer, pz.nx.zeros({"foo": 5, "bar": 7, "baz": 10}) 36 | ) 37 | 38 | 39 | if __name__ == "__main__": 40 | absltest.main() 41 | -------------------------------------------------------------------------------- /.github/workflows/unittests.yml: -------------------------------------------------------------------------------- 1 | name: Unittests 2 | 3 | on: 4 | # Allow to trigger the workflow manually (e.g. when deps changes) 5 | workflow_dispatch: 6 | # Run on pushes to main 7 | push: 8 | branches: 9 | - main 10 | # Run on pull requests to main (including test branches) 11 | pull_request: 12 | branches: 13 | - main 14 | 15 | jobs: 16 | unittest-job: 17 | runs-on: ubuntu-latest 18 | timeout-minutes: 30 19 | 20 | concurrency: 21 | group: ${{ github.workflow }}-${{ github.event_name }}-${{ github.ref }}-${{ github.head_ref || 'none' }} 22 | cancel-in-progress: true 23 | 24 | steps: 25 | - uses: actions/checkout@v3 26 | 27 | # Install deps 28 | - uses: actions/setup-python@v4 29 | with: 30 | python-version: 3.10.14 31 | # Uncomment to cache of pip dependencies (if tests too slow) 32 | # cache: pip 33 | # cache-dependency-path: '**/pyproject.toml' 34 | 35 | - uses: astral-sh/setup-uv@v3 36 | with: 37 | version: "0.4.17" 38 | 39 | - name: Install dependencies 40 | run: | 41 | uv sync --locked --extra extras --extra dev 42 | 43 | # Check formatting 44 | - name: Check pyink formatting 45 | run: uv run pyink penzai tests --check 46 | 47 | - name: Run pylint 48 | run: uv run pylint penzai 49 | 50 | # Run tests 51 | - name: Run tests 52 | run: uv run python run_tests.py 53 | 54 | # Run typechecker 55 | - name: Run pytype 56 | run: uv run pytype --jobs auto penzai 57 | -------------------------------------------------------------------------------- /tests/treescope/fixtures/treescope_examples_fixture.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | """Defines some classes for use in rendering tests.""" 16 | 17 | from __future__ import annotations 18 | 19 | from typing import Any 20 | 21 | from penzai import pz 22 | from penzai.deprecated.v1 import pz as pz_v1 23 | 24 | 25 | @pz.pytree_dataclass 26 | class StructWithOneChild(pz.Struct): 27 | foo: Any 28 | 29 | 30 | @pz.pytree_dataclass 31 | class ExampleLayer(pz_v1.Layer): 32 | foo: int 33 | 34 | def __call__(self, value: int) -> int: 35 | return value + self.foo 36 | 37 | 38 | @pz.pytree_dataclass 39 | class LayerThatHoldsStuff(pz_v1.Layer): 40 | stuff: Any 41 | 42 | def input_structure(self): 43 | return {"input": pz.chk.ArraySpec(shape=(1, 2, 3))} 44 | 45 | def output_structure(self): 46 | return {"output": pz.chk.ArraySpec(named_shape={"foo": 5})} 47 | 48 | @pz_v1.checked_layer_call 49 | def __call__(self, value: int) -> int: 50 | return value 51 | -------------------------------------------------------------------------------- /docs/_templates/pzmodule.rst: -------------------------------------------------------------------------------- 1 | {{ name | escape | underline}} 2 | 3 | .. automodule:: {{ fullname }} 4 | 5 | {% block classes %} 6 | {% if classes %} 7 | .. rubric:: {{ _('Classes') }} 8 | 9 | .. autosummary:: 10 | :toctree: leaf 11 | :template: pzclass.rst 12 | {% for item in classes %} 13 | {{ item }} 14 | {%- endfor %} 15 | {% endif %} 16 | {% endblock %} 17 | 18 | {% block functions %} 19 | {% if functions %} 20 | .. rubric:: {{ _('Functions') }} 21 | 22 | .. autosummary:: 23 | :toctree: leaf 24 | :template: pzbase.rst 25 | {% for item in functions %} 26 | {{ item }} 27 | {%- endfor %} 28 | {% endif %} 29 | {% endblock %} 30 | 31 | {% block attributes %} 32 | {% if attributes %} 33 | .. rubric:: {{ _('Module Attributes') }} 34 | 35 | .. autosummary:: 36 | :toctree: leaf 37 | :template: pzdata.rst 38 | {% for item in attributes %} 39 | {{ item }} 40 | {%- endfor %} 41 | {% endif %} 42 | {% endblock %} 43 | 44 | {% block exceptions %} 45 | {% if exceptions %} 46 | .. rubric:: {{ _('Exceptions') }} 47 | 48 | .. autosummary:: 49 | :toctree: leaf 50 | :template: pzbase.rst 51 | {% for item in exceptions %} 52 | {{ item }} 53 | {%- endfor %} 54 | {% endif %} 55 | {% endblock %} 56 | 57 | {% block modules %} 58 | {% if modules %} 59 | .. rubric:: Submodules 60 | 61 | .. autosummary:: 62 | :toctree: 63 | :template: pzmodule.rst 64 | :recursive: 65 | {% for item in modules %} 66 | {{ item }} 67 | {%- endfor %} 68 | {% endif %} 69 | {% endblock %} 70 | -------------------------------------------------------------------------------- /docs/_templates/pzmodule_full.rst: -------------------------------------------------------------------------------- 1 | {{ fullname | escape | underline}} 2 | 3 | .. automodule:: {{ fullname }} 4 | 5 | {% block classes %} 6 | {% if classes %} 7 | .. rubric:: {{ _('Classes') }} 8 | 9 | .. autosummary:: 10 | :toctree: leaf 11 | :template: pzclass.rst 12 | {% for item in classes %} 13 | {{ item }} 14 | {%- endfor %} 15 | {% endif %} 16 | {% endblock %} 17 | 18 | {% block functions %} 19 | {% if functions %} 20 | .. rubric:: {{ _('Functions') }} 21 | 22 | .. autosummary:: 23 | :toctree: leaf 24 | :template: pzbase.rst 25 | {% for item in functions %} 26 | {{ item }} 27 | {%- endfor %} 28 | {% endif %} 29 | {% endblock %} 30 | 31 | {% block attributes %} 32 | {% if attributes %} 33 | .. rubric:: {{ _('Module Attributes') }} 34 | 35 | .. autosummary:: 36 | :toctree: leaf 37 | :template: pzdata.rst 38 | {% for item in attributes %} 39 | {{ item }} 40 | {%- endfor %} 41 | {% endif %} 42 | {% endblock %} 43 | 44 | {% block exceptions %} 45 | {% if exceptions %} 46 | .. rubric:: {{ _('Exceptions') }} 47 | 48 | .. autosummary:: 49 | :toctree: leaf 50 | :template: pzbase.rst 51 | {% for item in exceptions %} 52 | {{ item }} 53 | {%- endfor %} 54 | {% endif %} 55 | {% endblock %} 56 | 57 | {% block modules %} 58 | {% if modules %} 59 | .. rubric:: Submodules 60 | 61 | .. autosummary:: 62 | :toctree: 63 | :template: pzmodule.rst 64 | :recursive: 65 | {% for item in modules %} 66 | {{ item }} 67 | {%- endfor %} 68 | {% endif %} 69 | {% endblock %} 70 | -------------------------------------------------------------------------------- /docs/ext/pz_alias_rewrite.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """Rewrites aliases in the pz namespace to point to their qualified versions.""" 15 | 16 | from typing import Any, Mapping 17 | 18 | from sphinx import application 19 | from treescope import canonical_aliases 20 | 21 | 22 | def substitute_pz_in_autodoc_docstring(app, what, name, obj, options, lines): 23 | """Rewrites docstrings for objects defined in `penzai.pz` to add alias links.""" 24 | del app, what, options 25 | if ( 26 | name.startswith("penzai.pz") 27 | or name.startswith("penzai.experimental.v2.pz") 28 | or name.startswith("penzai.deprecated.v1.pz") 29 | ): 30 | if not lines: 31 | lines.append("") 32 | alias = canonical_aliases.lookup_alias(obj) 33 | if alias is not None: 34 | lines[0] = f"*Alias of* :obj:`{str(alias)}`: " + lines[0] 35 | 36 | 37 | def setup(app: application.Sphinx) -> Mapping[str, Any]: 38 | app.connect("autodoc-process-docstring", substitute_pz_in_autodoc_docstring) 39 | return dict(version="0.1", parallel_read_safe=True) 40 | -------------------------------------------------------------------------------- /penzai/pz/ts.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | """Common aliases for treescope.""" 16 | 17 | # pylint: disable=g-multiple-import,g-importing-member,unused-import 18 | 19 | from penzai.treescope._compatibility_setup import ( 20 | display, 21 | register_as_default, 22 | basic_interactive_setup, 23 | ) 24 | from treescope import ( 25 | ArrayAutovisualizer, 26 | active_autovisualizer, 27 | active_expansion_strategy, 28 | active_renderer, 29 | Autovisualizer, 30 | ChildAutovisualizer, 31 | default_diverging_colormap, 32 | default_magic_autovisualizer, 33 | default_sequential_colormap, 34 | integer_digitbox, 35 | IPythonVisualization, 36 | register_autovisualize_magic, 37 | register_context_manager_magic, 38 | render_array_sharding, 39 | render_array, 40 | render_to_html, 41 | render_to_text, 42 | using_expansion_strategy, 43 | ) 44 | from treescope.figures import ( 45 | inline, 46 | indented, 47 | with_font_size, 48 | with_color, 49 | bolded, 50 | styled, 51 | text_on_color, 52 | ) 53 | 54 | vocab_autovisualizer = ArrayAutovisualizer.for_tokenizer 55 | -------------------------------------------------------------------------------- /tests/nn/embedding_test.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | """Tests for embedding layer.""" 16 | 17 | from absl.testing import absltest 18 | import jax 19 | from penzai import pz 20 | 21 | 22 | class EmbeddingTest(absltest.TestCase): 23 | 24 | def test_embeddings(self): 25 | table = pz.nn.EmbeddingTable.from_config( 26 | name="embeddings", 27 | init_base_rng=jax.random.key(2), 28 | vocab_size=17, 29 | embedding_axes={"foo": 5, "bar": 7}, 30 | vocabulary_axis="vocabulary", 31 | ) 32 | 33 | with self.subTest("lookup"): 34 | lookup_layer = pz.nn.EmbeddingLookup(table) 35 | result = lookup_layer(pz.nx.arange("baz", 3)) 36 | pz.chk.check_structure( 37 | result, pz.chk.ArraySpec(named_shape={"foo": 5, "bar": 7, "baz": 3}) 38 | ) 39 | 40 | with self.subTest("decode"): 41 | decode_layer = pz.nn.EmbeddingDecode(table) 42 | result = decode_layer(pz.nx.ones({"foo": 5, "bar": 7, "baz": 3})) 43 | pz.chk.check_structure( 44 | result, pz.chk.ArraySpec(named_shape={"baz": 3, "vocabulary": 17}) 45 | ) 46 | 47 | 48 | if __name__ == "__main__": 49 | absltest.main() 50 | -------------------------------------------------------------------------------- /tests/core/auto_order_types_test.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """Tests for auto_order_types.""" 15 | 16 | import dataclasses 17 | from absl.testing import absltest 18 | import jax 19 | from penzai.core import auto_order_types 20 | 21 | 22 | class AutoOrderTypesTest(absltest.TestCase): 23 | 24 | def test_auto_order_types(self): 25 | 26 | @dataclasses.dataclass(frozen=True) 27 | class Foo(auto_order_types.AutoOrderedAcrossTypes): 28 | a: int 29 | b: int 30 | 31 | @dataclasses.dataclass(frozen=True) 32 | class Bar(auto_order_types.AutoOrderedAcrossTypes): 33 | name: str 34 | 35 | items = [ 36 | Foo(a=1, b=2), 37 | Bar(name="bar"), 38 | Foo(a=3, b=4), 39 | Bar(name="bar2"), 40 | "some_string", 41 | Foo(a=5, b=6), 42 | Bar(name="bar3"), 43 | ] 44 | 45 | # Check that we can sort the items. 46 | _ = sorted(items) 47 | 48 | # Check that they can be used as keys in a dict and passed through JAX 49 | # pytree operations. 50 | my_dict = {val: i for i, val in enumerate(items)} 51 | leaves = jax.tree_util.tree_leaves(my_dict) 52 | self.assertSameElements(leaves, range(len(items))) 53 | 54 | 55 | if __name__ == "__main__": 56 | absltest.main() 57 | -------------------------------------------------------------------------------- /tests/nn/standardization_test.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | """Tests for standardization layers.""" 16 | 17 | from absl.testing import absltest 18 | import jax 19 | from penzai import pz 20 | 21 | 22 | class StandardizationTest(absltest.TestCase): 23 | 24 | def test_standardize_one(self): 25 | layer = pz.nn.Standardize(across="foo") 26 | result = layer(pz.nx.ones({"batch": 3, "foo": 4})) 27 | pz.chk.check_structure( 28 | result, pz.chk.ArraySpec(named_shape={"batch": 3, "foo": 4}) 29 | ) 30 | 31 | def test_standardize_multi(self): 32 | layer = pz.nn.Standardize(across=("foo", "bar")) 33 | result = layer(pz.nx.ones({"batch": 3, "foo": 4, "bar": 2})) 34 | pz.chk.check_structure( 35 | result, pz.chk.ArraySpec(named_shape={"batch": 3, "foo": 4, "bar": 2}) 36 | ) 37 | 38 | def test_layernorm(self): 39 | layer = pz.nn.LayerNorm.from_config( 40 | name="test", 41 | init_base_rng=jax.random.key(2), 42 | across_axes={"foo": 4, "bar": 2}, 43 | ) 44 | result = layer(pz.nx.ones({"batch": 3, "foo": 4, "bar": 2})) 45 | pz.chk.check_structure( 46 | result, pz.chk.ArraySpec(named_shape={"batch": 3, "foo": 4, "bar": 2}) 47 | ) 48 | 49 | 50 | if __name__ == "__main__": 51 | absltest.main() 52 | -------------------------------------------------------------------------------- /tests/core/misc_util_test.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | """Tests for miscellaneous self-contained utility objects.""" 16 | 17 | from absl.testing import absltest 18 | import jax 19 | from penzai import pz 20 | from penzai.core import tree_util 21 | 22 | 23 | class TreeUtilTest(absltest.TestCase): 24 | """Tests for core/tree_util.py.""" 25 | 26 | def test_tree_flatten_exactly_one_level_nested_pytree(self): 27 | result = tree_util.tree_flatten_exactly_one_level( 28 | {"a": [1, 2], "b": [3, 4]} 29 | ) 30 | self.assertIsNotNone(result) 31 | keys_and_subtrees, treedef = result 32 | self.assertEqual( 33 | keys_and_subtrees, 34 | [ 35 | (jax.tree_util.DictKey("a"), [1, 2]), 36 | (jax.tree_util.DictKey("b"), [3, 4]), 37 | ], 38 | ) 39 | self.assertEqual(treedef.unflatten(["x", "y"]), {"a": "x", "b": "y"}) 40 | 41 | def test_tree_flatten_exactly_one_level_leaf(self): 42 | self.assertIsNone(tree_util.tree_flatten_exactly_one_level(42)) 43 | 44 | 45 | class SliceLikeTest(absltest.TestCase): 46 | 47 | def test_slice(self): 48 | indexer = pz.slice[1, 2::10, :3, "foo"] 49 | self.assertEqual(indexer, (1, slice(2, None, 10), slice(None, 3), "foo")) 50 | 51 | 52 | if __name__ == "__main__": 53 | absltest.main() 54 | -------------------------------------------------------------------------------- /docs/api/pz.ts.rst: -------------------------------------------------------------------------------- 1 | ``pz.ts``: Treescope alias namespace 2 | ==================================== 3 | 4 | .. module:: penzai.pz.ts 5 | .. currentmodule:: penzai 6 | 7 | Treescope, Penzai's interactive pretty-printer, has moved to an independent 8 | package `treescope`. See the 9 | `Treescope documentation `_ 10 | for more information. 11 | 12 | The alias namespace ``pz.ts`` contains shorthand aliases for commonly-used 13 | Treescope functions. However, we recommend that new users use `treescope` 14 | directly. 15 | 16 | Using Treescope in IPython Notebooks 17 | ------------------------------------ 18 | 19 | .. autosummary:: 20 | 21 | pz.ts.basic_interactive_setup 22 | pz.ts.register_as_default 23 | pz.ts.register_autovisualize_magic 24 | pz.ts.register_context_manager_magic 25 | 26 | 27 | Showing Objects Explicitly 28 | -------------------------- 29 | 30 | .. autosummary:: 31 | pz.ts.render_array 32 | pz.ts.render_array_sharding 33 | pz.ts.integer_digitbox 34 | pz.ts.text_on_color 35 | pz.ts.display 36 | pz.show 37 | 38 | 39 | Styling Displayed Objects 40 | ------------------------- 41 | 42 | .. autosummary:: 43 | 44 | pz.ts.inline 45 | pz.ts.indented 46 | pz.ts.with_font_size 47 | pz.ts.with_color 48 | pz.ts.bolded 49 | pz.ts.styled 50 | 51 | 52 | Configuring Treescope 53 | --------------------- 54 | 55 | .. autosummary:: 56 | 57 | pz.ts.active_renderer 58 | pz.ts.active_expansion_strategy 59 | pz.ts.using_expansion_strategy 60 | pz.ts.active_autovisualizer 61 | pz.ts.default_diverging_colormap 62 | pz.ts.default_sequential_colormap 63 | 64 | Building Autovisualizers 65 | ------------------------ 66 | 67 | .. autosummary:: 68 | 69 | pz.ts.ArrayAutovisualizer 70 | pz.ts.Autovisualizer 71 | pz.ts.ChildAutovisualizer 72 | pz.ts.IPythonVisualization 73 | pz.ts.vocab_autovisualizer 74 | pz.ts.default_magic_autovisualizer 75 | 76 | 77 | Rendering to Strings 78 | -------------------- 79 | 80 | .. autosummary:: 81 | 82 | pz.ts.render_to_text 83 | pz.ts.render_to_html 84 | 85 | -------------------------------------------------------------------------------- /penzai/core/random_stream.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | """A stateful random stream helper.""" 16 | 17 | from __future__ import annotations 18 | 19 | import jax 20 | from penzai.core import struct 21 | from penzai.core import variables 22 | 23 | 24 | @struct.pytree_dataclass 25 | class RandomStream(struct.Struct): 26 | """A stateful random stream object. 27 | 28 | This object can be used to generate a sequence of random numbers inside a 29 | Penzai model. It uses a Variable to track its state. 30 | 31 | Attributes: 32 | base_key: The base key to use for this random stream. This does not change, 33 | and determines which sequence of random numbers will be generated. 34 | offset: The number of random numbers that have been generated so far. 35 | """ 36 | 37 | base_key: jax.Array 38 | offset: variables.StateVariable[int | jax.Array] 39 | 40 | @classmethod 41 | def from_base_key( 42 | cls, 43 | base_key: jax.Array, 44 | offset_label: variables.VariableLabel = "random_stream_offset", 45 | ) -> RandomStream: 46 | """Returns a new random stream with the given base key.""" 47 | return cls( 48 | base_key=base_key, offset=variables.StateVariable(0, label=offset_label) 49 | ) 50 | 51 | def next_key(self) -> jax.Array: 52 | """Returns the next key in the sequence, and advances the stream.""" 53 | old_offset = self.offset.value 54 | self.offset.value = self.offset.value + 1 55 | return jax.random.fold_in(self.base_key, old_offset) 56 | -------------------------------------------------------------------------------- /docs/_templates/pzclass.rst: -------------------------------------------------------------------------------- 1 | {{ name | escape | underline}} 2 | 3 | .. currentmodule:: {{ module }} 4 | 5 | .. autoclass:: {{ objname }} 6 | :members: 7 | :special-members: 8 | :show-inheritance: 9 | 10 | 11 | {% set attr_ns = namespace(inherited=false) %} 12 | {% for item in attributes %} 13 | {% if item in inherited_members %} 14 | {% set attr_ns.inherited = true %} 15 | {% endif %} 16 | {%- endfor %} 17 | 18 | {% if attr_ns.inherited %} 19 | .. rubric:: {{ _('Inherited Attributes') }} 20 | .. autosummary:: 21 | {% for item in attributes %} 22 | {% if item in inherited_members %} 23 | ~{{ name }}.{{ item }} 24 | {% endif %} 25 | {%- endfor %} 26 | {% endif %} 27 | 28 | {% set method_ns = namespace(methods_ext=methods, own=false, inherited=false) %} 29 | {% for special in ["__call__", "__enter__", "__exit__"] %} 30 | {% if special in members %} 31 | {% set method_ns.methods_ext = method_ns.methods_ext + [special] %} 32 | {% endif %} 33 | {%- endfor %} 34 | {% for item in method_ns.methods_ext %} 35 | {% if item in inherited_members %} 36 | {% set method_ns.inherited = true %} 37 | {% else %} 38 | {% set method_ns.own = true %} 39 | {% endif %} 40 | {%- endfor %} 41 | 42 | {% if method_ns.own %} 43 | .. rubric:: {{ _('Methods') }} 44 | .. autosummary:: 45 | {% for item in method_ns.methods_ext %} 46 | {% if item not in inherited_members %} 47 | ~{{ name }}.{{ item }} 48 | {% endif %} 49 | {%- endfor %} 50 | {% endif %} 51 | 52 | {% if attributes %} 53 | .. rubric:: {{ _('Attributes') }} 54 | .. autosummary:: 55 | {% for item in attributes %} 56 | ~{{ name }}.{{ item }} 57 | {%- endfor %} 58 | {% endif %} 59 | 60 | {% if method_ns.inherited %} 61 | 62 | .. rubric:: {{ _('Inherited Methods') }} 63 | .. raw:: html 64 | 65 |
66 | (expand to view inherited methods) 67 | 68 | .. autosummary:: 69 | {% for item in method_ns.methods_ext %} 70 | {% if item in inherited_members %} 71 | ~{{ name }}.{{ item }} 72 | {% endif %} 73 | {%- endfor %} 74 | 75 | .. raw:: html 76 | 77 |
78 | 79 | {% endif %} 80 | -------------------------------------------------------------------------------- /docs/ext/nb_output_cell_to_iframe.py: -------------------------------------------------------------------------------- 1 | # Copyright 2024 The Penzai Authors. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """Wraps MyST-nb output cells in HTML templates. 15 | 16 | This can be used to sandbox the environment of output cell renderings, which is 17 | useful for Treescope renderings since they make heavy use of Javascript and CSS, 18 | and can also improve responsiveness of the main page. 19 | 20 | This transformation replaces output cells with ``