tensorgami 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (74) hide show
  1. tensorgami-0.1.0/.gitignore +20 -0
  2. tensorgami-0.1.0/CHANGELOG.md +54 -0
  3. tensorgami-0.1.0/CITATION.cff +16 -0
  4. tensorgami-0.1.0/LICENSE +21 -0
  5. tensorgami-0.1.0/PKG-INFO +172 -0
  6. tensorgami-0.1.0/README.md +140 -0
  7. tensorgami-0.1.0/THIRD_PARTY_NOTICES.md +112 -0
  8. tensorgami-0.1.0/pyproject.toml +117 -0
  9. tensorgami-0.1.0/src/tensorgami/__init__.py +22 -0
  10. tensorgami-0.1.0/src/tensorgami/api/__init__.py +9 -0
  11. tensorgami-0.1.0/src/tensorgami/api/annotations.py +43 -0
  12. tensorgami-0.1.0/src/tensorgami/api/draw.py +187 -0
  13. tensorgami-0.1.0/src/tensorgami/api/figure.py +119 -0
  14. tensorgami-0.1.0/src/tensorgami/importers/__init__.py +13 -0
  15. tensorgami-0.1.0/src/tensorgami/importers/base.py +21 -0
  16. tensorgami-0.1.0/src/tensorgami/importers/torch/__init__.py +15 -0
  17. tensorgami-0.1.0/src/tensorgami/importers/torch/catalog.py +215 -0
  18. tensorgami-0.1.0/src/tensorgami/importers/torch/convert.py +193 -0
  19. tensorgami-0.1.0/src/tensorgami/importers/torch/importer.py +207 -0
  20. tensorgami-0.1.0/src/tensorgami/importers/torch/isolation.py +39 -0
  21. tensorgami-0.1.0/src/tensorgami/importers/torch/overrides.py +76 -0
  22. tensorgami-0.1.0/src/tensorgami/importers/torch/recorder.py +219 -0
  23. tensorgami-0.1.0/src/tensorgami/ir/__init__.py +74 -0
  24. tensorgami-0.1.0/src/tensorgami/ir/builder.py +182 -0
  25. tensorgami-0.1.0/src/tensorgami/ir/model.py +308 -0
  26. tensorgami-0.1.0/src/tensorgami/ir/serialize.py +153 -0
  27. tensorgami-0.1.0/src/tensorgami/ir/shapes.py +21 -0
  28. tensorgami-0.1.0/src/tensorgami/ir/summary.py +43 -0
  29. tensorgami-0.1.0/src/tensorgami/ir/validation.py +228 -0
  30. tensorgami-0.1.0/src/tensorgami/layout/__init__.py +27 -0
  31. tensorgami-0.1.0/src/tensorgami/layout/base.py +85 -0
  32. tensorgami-0.1.0/src/tensorgami/layout/cell.py +93 -0
  33. tensorgami-0.1.0/src/tensorgami/layout/labels.py +270 -0
  34. tensorgami-0.1.0/src/tensorgami/layout/layered.py +1220 -0
  35. tensorgami-0.1.0/src/tensorgami/layout/unrolled.py +221 -0
  36. tensorgami-0.1.0/src/tensorgami/layout/ushape.py +196 -0
  37. tensorgami-0.1.0/src/tensorgami/layout/volume.py +34 -0
  38. tensorgami-0.1.0/src/tensorgami/render/__init__.py +31 -0
  39. tensorgami-0.1.0/src/tensorgami/render/canvas.py +56 -0
  40. tensorgami-0.1.0/src/tensorgami/render/edges.py +180 -0
  41. tensorgami-0.1.0/src/tensorgami/render/groups.py +55 -0
  42. tensorgami-0.1.0/src/tensorgami/render/legend.py +85 -0
  43. tensorgami-0.1.0/src/tensorgami/render/mpl.py +242 -0
  44. tensorgami-0.1.0/src/tensorgami/render/painters.py +315 -0
  45. tensorgami-0.1.0/src/tensorgami/render/scene_renderer.py +79 -0
  46. tensorgami-0.1.0/src/tensorgami/render/svg.py +145 -0
  47. tensorgami-0.1.0/src/tensorgami/render/text.py +97 -0
  48. tensorgami-0.1.0/src/tensorgami/scene.py +261 -0
  49. tensorgami-0.1.0/src/tensorgami/style/__init__.py +73 -0
  50. tensorgami-0.1.0/src/tensorgami/style/color.py +109 -0
  51. tensorgami-0.1.0/src/tensorgami/style/glyphs.py +148 -0
  52. tensorgami-0.1.0/src/tensorgami/style/loader.py +265 -0
  53. tensorgami-0.1.0/src/tensorgami/style/theme.py +65 -0
  54. tensorgami-0.1.0/src/tensorgami/style/themes/__init__.py +1 -0
  55. tensorgami-0.1.0/src/tensorgami/style/themes/default.toml +237 -0
  56. tensorgami-0.1.0/src/tensorgami/style/themes/origami-extended.toml +259 -0
  57. tensorgami-0.1.0/src/tensorgami/style/themes/origami.toml +263 -0
  58. tensorgami-0.1.0/src/tensorgami/style/validation.py +139 -0
  59. tensorgami-0.1.0/src/tensorgami/transforms/__init__.py +37 -0
  60. tensorgami-0.1.0/src/tensorgami/transforms/annotate.py +229 -0
  61. tensorgami-0.1.0/src/tensorgami/transforms/base.py +185 -0
  62. tensorgami-0.1.0/src/tensorgami/transforms/cells.py +340 -0
  63. tensorgami-0.1.0/src/tensorgami/transforms/cleanup.py +198 -0
  64. tensorgami-0.1.0/src/tensorgami/transforms/collapse.py +139 -0
  65. tensorgami-0.1.0/src/tensorgami/transforms/levels.py +74 -0
  66. tensorgami-0.1.0/src/tensorgami/transforms/recurrence.py +58 -0
  67. tensorgami-0.1.0/src/tensorgami/transforms/repeats.py +174 -0
  68. tensorgami-0.1.0/src/tensorgami/transforms/unroll.py +342 -0
  69. tensorgami-0.1.0/src/tensorgami/transforms/ushape.py +147 -0
  70. tensorgami-0.1.0/src/tensorgami/transforms/volume.py +40 -0
  71. tensorgami-0.1.0/src/tensorgami/views/__init__.py +55 -0
  72. tensorgami-0.1.0/src/tensorgami/views/base.py +25 -0
  73. tensorgami-0.1.0/src/tensorgami/views/block.py +125 -0
  74. tensorgami-0.1.0/src/tensorgami/vocabulary.py +43 -0
@@ -0,0 +1,20 @@
1
+ # Python
2
+ __pycache__/
3
+ *.py[cod]
4
+ *.egg-info/
5
+ build/
6
+ dist/
7
+
8
+ # Environments
9
+ .venv/
10
+
11
+ # Tooling caches
12
+ .pytest_cache/
13
+ .ruff_cache/
14
+ .coverage
15
+ htmlcov/
16
+
17
+ # OS / editors
18
+ .DS_Store
19
+ .vscode/
20
+ .idea/
@@ -0,0 +1,54 @@
1
+ # Changelog
2
+
3
+ All notable changes to Tensorgami. The format follows
4
+ [Keep a Changelog](https://keepachangelog.com/en/1.1.0/); versions follow
5
+ [Semantic Versioning](https://semver.org/) (0.x: the API may still change between minor
6
+ versions).
7
+
8
+ ## [Unreleased]
9
+
10
+ ## [0.1.0] — 2026-10-13
11
+
12
+ The first release: a PyTorch model to a figure in one call.
13
+
14
+ ### Drawing
15
+
16
+ - `tg.draw(model, example_input, …)` traces a model once (no gradients, eval mode; weights,
17
+ training flags and the random state are left as they were) and returns a `tg.Figure`: shown
18
+ inline in Jupyter, saved on request as PNG, SVG, PDF, JPG or TIFF (`fig.save`), drawn into
19
+ a panel of your own matplotlib figure (`fig.draw_on(ax)`).
20
+ - Levels of detail: `level=0` (system) … `3` (every operation), by number or name.
21
+ - Views: `block` (boxes and wires), `unrolled` (recurrence as a time × layer grid),
22
+ `volume` (CNN layers as tensor volumes), `cell` (an LSTM / GRU cell as its gates), and
23
+ `auto` (unrolled for traced loops over time, block otherwise).
24
+ - Encoder–decoders with nested skips (U-Net, pix2pix) fold into a U; `u_levels` limits how
25
+ many levels are drawn.
26
+ - Figures wider than `wrap_width` (1100 px by default) fold into rows that read like lines of
27
+ text, frames included: a module too wide for the page wraps inside its frame. A residual
28
+ block is never split.
29
+ - Several models in one figure (`tg.draw({"generator": G, "discriminator": D}, …,
30
+ links=[…])`) and annotations: `tg.Loss` (with training arrows), `tg.Input`, `tg.Note`.
31
+ - Models on PyTorch's `meta` device (`model.to("meta")`) are traced without allocating
32
+ memory for weights or activations (#2, @HarrisonKramer).
33
+ - `overrides` set the label, category or leaf-ness of a module; clear errors before the
34
+ model runs (misspelt options, unknown styles, levels, views, references).
35
+
36
+ ### Styles
37
+
38
+ - `default`, `origami` (upright cards left → right) and `origami-extended` (flat cards, top →
39
+ bottom); a style is one TOML file, and your own can be passed by path.
40
+
41
+ ### Examples
42
+
43
+ - Notebooks, already run: AlexNet, ResNet-18/34/50, pix2pix, and a tour of one model per
44
+ family; `examples/alexnet.py` and `examples/resnet.py` define the models as torchvision
45
+ does, without needing it.
46
+
47
+ ### Under the hood
48
+
49
+ - Importer (PyTorch tracing), framework-neutral IR, transforms, layered layout with block
50
+ layouts (unrolled grid, U-shape, LSTM cell), SVG and matplotlib renderers.
51
+ - An acceptance gallery (`docs/gallery/`) of every model family in every view and style,
52
+ checked against committed SVGs and layout invariants by the test suite.
53
+ - Development on Windows: `uv sync` (CPU PyTorch) and the tests work there too (#1,
54
+ @HarrisonKramer).
@@ -0,0 +1,16 @@
1
+ cff-version: 1.2.0
2
+ message: "If you use Tensorgami for figures in your work, please cite it as below."
3
+ title: "Tensorgami: draw the architecture and information flow of PyTorch models"
4
+ type: software
5
+ authors:
6
+ - family-names: "Utina"
7
+ given-names: "Andrei"
8
+ version: "0.1.0"
9
+ date-released: "2026-10-13"
10
+ license: MIT
11
+ repository-code: "https://github.com/AndreiUtina/TensorGami"
12
+ keywords:
13
+ - pytorch
14
+ - neural networks
15
+ - architecture diagrams
16
+ - visualization
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Andrei Utina
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,172 @@
1
+ Metadata-Version: 2.5
2
+ Name: tensorgami
3
+ Version: 0.1.0
4
+ Summary: Draw the architecture and information flow of PyTorch models: figures for papers, reports and slides
5
+ Project-URL: Homepage, https://github.com/AndreiUtina/TensorGami
6
+ Project-URL: Repository, https://github.com/AndreiUtina/TensorGami
7
+ Project-URL: Issues, https://github.com/AndreiUtina/TensorGami/issues
8
+ Project-URL: Changelog, https://github.com/AndreiUtina/TensorGami/blob/main/CHANGELOG.md
9
+ Project-URL: Gallery, https://github.com/AndreiUtina/TensorGami/blob/main/docs/gallery/README.md
10
+ Author: Andrei Utina
11
+ License-Expression: MIT
12
+ License-File: LICENSE
13
+ Keywords: architecture,deep learning,diagram,figure,matplotlib,neural network,pytorch,visualization
14
+ Classifier: Development Status :: 3 - Alpha
15
+ Classifier: Framework :: Jupyter
16
+ Classifier: Intended Audience :: Education
17
+ Classifier: Intended Audience :: Science/Research
18
+ Classifier: Operating System :: OS Independent
19
+ Classifier: Programming Language :: Python :: 3
20
+ Classifier: Programming Language :: Python :: 3 :: Only
21
+ Classifier: Programming Language :: Python :: 3.11
22
+ Classifier: Programming Language :: Python :: 3.12
23
+ Classifier: Programming Language :: Python :: 3.13
24
+ Classifier: Topic :: Multimedia :: Graphics :: Presentation
25
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
26
+ Classifier: Topic :: Scientific/Engineering :: Visualization
27
+ Requires-Python: >=3.11
28
+ Requires-Dist: matplotlib>=3.7
29
+ Provides-Extra: torch
30
+ Requires-Dist: torch>=2.2; extra == 'torch'
31
+ Description-Content-Type: text/markdown
32
+
33
+ ![TensorGami](https://raw.githubusercontent.com/AndreiUtina/TensorGami/main/docs/brand/logo-banner.png)
34
+
35
+ # Tensorgami
36
+
37
+ [![CI](https://github.com/AndreiUtina/TensorGami/actions/workflows/ci.yml/badge.svg)](https://github.com/AndreiUtina/TensorGami/actions/workflows/ci.yml)
38
+ [![PyPI](https://img.shields.io/pypi/v/tensorgami)](https://pypi.org/project/tensorgami/)
39
+ [![Python](https://img.shields.io/pypi/pyversions/tensorgami)](https://pypi.org/project/tensorgami/)
40
+ [![License: MIT](https://img.shields.io/badge/license-MIT-blue)](https://github.com/AndreiUtina/TensorGami/blob/main/LICENSE)
41
+
42
+ Draw the architecture and information flow of PyTorch models, locally, in Python: figures
43
+ for papers, reports and slides. Call `tg.draw` on a loaded model in a script or a notebook,
44
+ see the figure inline in Jupyter, and save it as PNG, SVG or PDF.
45
+
46
+ ![AlexNet in the origami style](https://raw.githubusercontent.com/AndreiUtina/TensorGami/main/docs/gallery/origami/alexnet.svg)
47
+
48
+ ![pix2pix: U-Net generator and PatchGAN discriminator as tensor volumes](https://raw.githubusercontent.com/AndreiUtina/TensorGami/main/docs/gallery/pix2pix-system.svg)
49
+
50
+ ## Install
51
+
52
+ ```sh
53
+ pip install tensorgami
54
+ ```
55
+
56
+ Tensorgami needs Python ≥ 3.11 and matplotlib. To draw a model it runs the model once with
57
+ PyTorch (≥ 2.2), which it does not install itself, so your CPU or CUDA build stays as it is.
58
+ If PyTorch isn't installed yet, install it first
59
+ ([pytorch.org](https://pytorch.org/get-started/locally/)), or let pip pick the default build
60
+ with `pip install "tensorgami[torch]"`.
61
+
62
+ The latest development version: `pip install "git+https://github.com/AndreiUtina/TensorGami"`.
63
+
64
+ ## Quick start
65
+
66
+ ```python
67
+ import torch
68
+ from torch import nn
69
+ import tensorgami as tg
70
+
71
+ model = nn.Sequential(
72
+ nn.Conv2d(3, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
73
+ nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(), nn.AdaptiveAvgPool2d(1),
74
+ nn.Flatten(), nn.Linear(32, 10),
75
+ )
76
+
77
+ fig = tg.draw(model, torch.randn(1, 3, 32, 32)) # runs the model once, leaves it untouched
78
+ fig # shows inline in Jupyter
79
+ fig.save("model.png", dpi=300) # or .svg / .pdf; nothing is saved before
80
+ ```
81
+
82
+ The example input is what the model is called with: a tensor, a tuple of positional
83
+ arguments, or a dict of keyword arguments, including the batch dimension.
84
+
85
+ ## What you can draw
86
+
87
+ ```python
88
+ tg.draw(model, x, level="stages") # less detail: 0 system … 3 operators
89
+ tg.draw(model, x, style="origami") # or "origami-extended", or a path to a TOML
90
+ tg.draw(model, x, direction="TB", legend=False) # top → bottom, no legend
91
+ tg.draw(model, x, wrap_width=1600) # fold into rows beyond 1600 px (None: never)
92
+ tg.draw(model, x, view="unrolled") # recurrence as a time × layer grid
93
+ tg.draw(model, x, view="volume") # CNN layers as tensor volumes
94
+ tg.draw(model, x, view="cell") # an LSTM / GRU cell as its gates
95
+ tg.draw(unet, x, u_shape=False) # keep a U-Net on one line (default: a U)
96
+ tg.draw(model, x, overrides={"classifier.6": {"label": "logits"}}) # rename a box
97
+ tg.draw(vae, x, annotations=[tg.Loss("KL divergence", "mu", "logvar")]) # a loss term
98
+ tg.draw({"generator": G, "discriminator": D}, {"generator": (z, y), "discriminator": (x, y)},
99
+ links=[("generator", "discriminator.input 0")], # several models, wired together
100
+ annotations=[tg.Loss("adversarial", "discriminator", trains=["generator", "discriminator"])])
101
+ ```
102
+
103
+ - **Levels of detail**: the whole system, stages, layers or single operators. Repeated blocks
104
+ fold into one frame marked ×N, activations fold into the layer before them, and shapes
105
+ appear on the wires where they change.
106
+ - **Views**: block diagrams (`view="block"`), recurrent networks unrolled over time
107
+ (`"unrolled"`), CNNs as 3-D tensor volumes (`"volume"`) and recurrent cells as their gates
108
+ (`"cell"`). Encoder–decoders with nested skips (U-Net, pix2pix) fold into a U.
109
+ - **Styles**: [`default`](https://github.com/AndreiUtina/TensorGami/blob/main/docs/style/style-baseline.md),
110
+ and [`origami` and `origami-extended`](https://github.com/AndreiUtina/TensorGami/blob/main/docs/style/origami.md)
111
+ (narrow boxes in the style of the survey figures). A style is a TOML file; your own can be
112
+ passed by path.
113
+ - **Several models in one figure**: GANs, VAEs and other training set-ups, with links between
114
+ the models, extra inputs, loss terms with dashed training arrows, and notes.
115
+
116
+ Model families covered by the tests: CNN (AlexNet, ResNet), LSTM, ConvLSTM, cGAN, VAE,
117
+ Transformer and pix2pix. The [acceptance gallery](https://github.com/AndreiUtina/TensorGami/blob/main/docs/gallery/README.md)
118
+ shows each of them in every view and style. Every figure there is drawn automatically, and
119
+ the tests check it.
120
+
121
+ ## Examples
122
+
123
+ Notebooks with every option, already run, are in
124
+ [`examples/`](https://github.com/AndreiUtina/TensorGami/blob/main/examples):
125
+
126
+ - [`alexnet.ipynb`](https://github.com/AndreiUtina/TensorGami/blob/main/examples/alexnet.ipynb):
127
+ levels, views, styles, layout options, saving, panels;
128
+ - [`resnet.ipynb`](https://github.com/AndreiUtina/TensorGami/blob/main/examples/resnet.ipynb):
129
+ ResNet-18/34/50, how to read the figure, wide figures folded into rows;
130
+ - [`pix2pix.ipynb`](https://github.com/AndreiUtina/TensorGami/blob/main/examples/pix2pix.ipynb):
131
+ a U-Net generator, a PatchGAN discriminator and both together;
132
+ - [`tour.ipynb`](https://github.com/AndreiUtina/TensorGami/blob/main/examples/tour.ipynb):
133
+ one model per family.
134
+
135
+ The pipeline's stages can also be used one by one, e.g. to inspect what was captured:
136
+
137
+ ```python
138
+ from tensorgami.importers.torch import trace
139
+ from tensorgami.transforms import simplify
140
+
141
+ graph = trace(model, torch.randn(1, 3, 32, 32))
142
+ print(graph.summary()) # every operation, its module and shape
143
+ print(simplify(graph, level=2).summary()) # layers: activations folded, repeats ×N
144
+ ```
145
+
146
+ Background: the
147
+ [literature survey](https://github.com/AndreiUtina/TensorGami/blob/main/docs/survey/tensorgami-literature-survey.pdf)
148
+ on how neural-network architectures are drawn, and the
149
+ [architecture](https://github.com/AndreiUtina/TensorGami/blob/main/docs/architecture.md)
150
+ of Tensorgami itself.
151
+
152
+ ## Development
153
+
154
+ ```sh
155
+ git clone https://github.com/AndreiUtina/TensorGami && cd TensorGami
156
+ uv sync # the environment, with CPU PyTorch
157
+ uv run pytest # the tests, including the gallery
158
+ uv run ruff check . && uv run ruff format . # lint / format
159
+ uv run python docs/gallery/build_gallery.py # accept changed drawings: rebuild the gallery
160
+ cd examples && uv run --with nbclient --with nbformat --with ipykernel jupyter execute --inplace *.ipynb
161
+ ```
162
+
163
+ Releases are described in [docs/RELEASING.md](https://github.com/AndreiUtina/TensorGami/blob/main/docs/RELEASING.md),
164
+ changes in the [changelog](https://github.com/AndreiUtina/TensorGami/blob/main/CHANGELOG.md).
165
+
166
+ ## License
167
+
168
+ MIT, see [LICENSE](https://github.com/AndreiUtina/TensorGami/blob/main/LICENSE). The example
169
+ and test models adapted from torchvision and pytorch-CycleGAN-and-pix2pix keep their BSD
170
+ licenses, see [THIRD_PARTY_NOTICES.md](https://github.com/AndreiUtina/TensorGami/blob/main/THIRD_PARTY_NOTICES.md).
171
+ To cite Tensorgami, use [CITATION.cff](https://github.com/AndreiUtina/TensorGami/blob/main/CITATION.cff)
172
+ (GitHub's "Cite this repository").
@@ -0,0 +1,140 @@
1
+ ![TensorGami](docs/brand/logo-banner.png)
2
+
3
+ # Tensorgami
4
+
5
+ [![CI](https://github.com/AndreiUtina/TensorGami/actions/workflows/ci.yml/badge.svg)](https://github.com/AndreiUtina/TensorGami/actions/workflows/ci.yml)
6
+ [![PyPI](https://img.shields.io/pypi/v/tensorgami)](https://pypi.org/project/tensorgami/)
7
+ [![Python](https://img.shields.io/pypi/pyversions/tensorgami)](https://pypi.org/project/tensorgami/)
8
+ [![License: MIT](https://img.shields.io/badge/license-MIT-blue)](LICENSE)
9
+
10
+ Draw the architecture and information flow of PyTorch models, locally, in Python: figures
11
+ for papers, reports and slides. Call `tg.draw` on a loaded model in a script or a notebook,
12
+ see the figure inline in Jupyter, and save it as PNG, SVG or PDF.
13
+
14
+ ![AlexNet in the origami style](docs/gallery/origami/alexnet.svg)
15
+
16
+ ![pix2pix: U-Net generator and PatchGAN discriminator as tensor volumes](docs/gallery/pix2pix-system.svg)
17
+
18
+ ## Install
19
+
20
+ ```sh
21
+ pip install tensorgami
22
+ ```
23
+
24
+ Tensorgami needs Python ≥ 3.11 and matplotlib. To draw a model it runs the model once with
25
+ PyTorch (≥ 2.2), which it does not install itself, so your CPU or CUDA build stays as it is.
26
+ If PyTorch isn't installed yet, install it first
27
+ ([pytorch.org](https://pytorch.org/get-started/locally/)), or let pip pick the default build
28
+ with `pip install "tensorgami[torch]"`.
29
+
30
+ The latest development version: `pip install "git+https://github.com/AndreiUtina/TensorGami"`.
31
+
32
+ ## Quick start
33
+
34
+ ```python
35
+ import torch
36
+ from torch import nn
37
+ import tensorgami as tg
38
+
39
+ model = nn.Sequential(
40
+ nn.Conv2d(3, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
41
+ nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(), nn.AdaptiveAvgPool2d(1),
42
+ nn.Flatten(), nn.Linear(32, 10),
43
+ )
44
+
45
+ fig = tg.draw(model, torch.randn(1, 3, 32, 32)) # runs the model once, leaves it untouched
46
+ fig # shows inline in Jupyter
47
+ fig.save("model.png", dpi=300) # or .svg / .pdf; nothing is saved before
48
+ ```
49
+
50
+ The example input is what the model is called with: a tensor, a tuple of positional
51
+ arguments, or a dict of keyword arguments, including the batch dimension.
52
+
53
+ ## What you can draw
54
+
55
+ ```python
56
+ tg.draw(model, x, level="stages") # less detail: 0 system … 3 operators
57
+ tg.draw(model, x, style="origami") # or "origami-extended", or a path to a TOML
58
+ tg.draw(model, x, direction="TB", legend=False) # top → bottom, no legend
59
+ tg.draw(model, x, wrap_width=1600) # fold into rows beyond 1600 px (None: never)
60
+ tg.draw(model, x, view="unrolled") # recurrence as a time × layer grid
61
+ tg.draw(model, x, view="volume") # CNN layers as tensor volumes
62
+ tg.draw(model, x, view="cell") # an LSTM / GRU cell as its gates
63
+ tg.draw(unet, x, u_shape=False) # keep a U-Net on one line (default: a U)
64
+ tg.draw(model, x, overrides={"classifier.6": {"label": "logits"}}) # rename a box
65
+ tg.draw(vae, x, annotations=[tg.Loss("KL divergence", "mu", "logvar")]) # a loss term
66
+ tg.draw({"generator": G, "discriminator": D}, {"generator": (z, y), "discriminator": (x, y)},
67
+ links=[("generator", "discriminator.input 0")], # several models, wired together
68
+ annotations=[tg.Loss("adversarial", "discriminator", trains=["generator", "discriminator"])])
69
+ ```
70
+
71
+ - **Levels of detail**: the whole system, stages, layers or single operators. Repeated blocks
72
+ fold into one frame marked ×N, activations fold into the layer before them, and shapes
73
+ appear on the wires where they change.
74
+ - **Views**: block diagrams (`view="block"`), recurrent networks unrolled over time
75
+ (`"unrolled"`), CNNs as 3-D tensor volumes (`"volume"`) and recurrent cells as their gates
76
+ (`"cell"`). Encoder–decoders with nested skips (U-Net, pix2pix) fold into a U.
77
+ - **Styles**: [`default`](docs/style/style-baseline.md),
78
+ and [`origami` and `origami-extended`](docs/style/origami.md)
79
+ (narrow boxes in the style of the survey figures). A style is a TOML file; your own can be
80
+ passed by path.
81
+ - **Several models in one figure**: GANs, VAEs and other training set-ups, with links between
82
+ the models, extra inputs, loss terms with dashed training arrows, and notes.
83
+
84
+ Model families covered by the tests: CNN (AlexNet, ResNet), LSTM, ConvLSTM, cGAN, VAE,
85
+ Transformer and pix2pix. The [acceptance gallery](docs/gallery/README.md)
86
+ shows each of them in every view and style. Every figure there is drawn automatically, and
87
+ the tests check it.
88
+
89
+ ## Examples
90
+
91
+ Notebooks with every option, already run, are in
92
+ [`examples/`](examples):
93
+
94
+ - [`alexnet.ipynb`](examples/alexnet.ipynb):
95
+ levels, views, styles, layout options, saving, panels;
96
+ - [`resnet.ipynb`](examples/resnet.ipynb):
97
+ ResNet-18/34/50, how to read the figure, wide figures folded into rows;
98
+ - [`pix2pix.ipynb`](examples/pix2pix.ipynb):
99
+ a U-Net generator, a PatchGAN discriminator and both together;
100
+ - [`tour.ipynb`](examples/tour.ipynb):
101
+ one model per family.
102
+
103
+ The pipeline's stages can also be used one by one, e.g. to inspect what was captured:
104
+
105
+ ```python
106
+ from tensorgami.importers.torch import trace
107
+ from tensorgami.transforms import simplify
108
+
109
+ graph = trace(model, torch.randn(1, 3, 32, 32))
110
+ print(graph.summary()) # every operation, its module and shape
111
+ print(simplify(graph, level=2).summary()) # layers: activations folded, repeats ×N
112
+ ```
113
+
114
+ Background: the
115
+ [literature survey](docs/survey/tensorgami-literature-survey.pdf)
116
+ on how neural-network architectures are drawn, and the
117
+ [architecture](docs/architecture.md)
118
+ of Tensorgami itself.
119
+
120
+ ## Development
121
+
122
+ ```sh
123
+ git clone https://github.com/AndreiUtina/TensorGami && cd TensorGami
124
+ uv sync # the environment, with CPU PyTorch
125
+ uv run pytest # the tests, including the gallery
126
+ uv run ruff check . && uv run ruff format . # lint / format
127
+ uv run python docs/gallery/build_gallery.py # accept changed drawings: rebuild the gallery
128
+ cd examples && uv run --with nbclient --with nbformat --with ipykernel jupyter execute --inplace *.ipynb
129
+ ```
130
+
131
+ Releases are described in [docs/RELEASING.md](docs/RELEASING.md),
132
+ changes in the [changelog](CHANGELOG.md).
133
+
134
+ ## License
135
+
136
+ MIT, see [LICENSE](LICENSE). The example
137
+ and test models adapted from torchvision and pytorch-CycleGAN-and-pix2pix keep their BSD
138
+ licenses, see [THIRD_PARTY_NOTICES.md](THIRD_PARTY_NOTICES.md).
139
+ To cite Tensorgami, use [CITATION.cff](CITATION.cff)
140
+ (GitHub's "Cite this repository").
@@ -0,0 +1,112 @@
1
+ # Third-party notices
2
+
3
+ Tensorgami itself is released under the [MIT License](LICENSE). The repository also contains
4
+ model definitions adapted from other projects, used as examples and test models. They are not
5
+ part of the installed package (`src/tensorgami`); their licenses are reproduced here, as those
6
+ licenses require.
7
+
8
+ ## AlexNet and ResNet — torchvision
9
+
10
+ `examples/alexnet.py` defines AlexNet exactly as `torchvision.models.alexnet` does, and
11
+ `examples/resnet.py` ResNet-18/34/50 as `torchvision.models.resnet` does (without its options
12
+ for grouped and dilated convolutions), from <https://github.com/pytorch/vision>, so that
13
+ torchvision is not needed.
14
+
15
+ ```text
16
+ BSD 3-Clause License
17
+
18
+ Copyright (c) Soumith Chintala 2016,
19
+ All rights reserved.
20
+
21
+ Redistribution and use in source and binary forms, with or without
22
+ modification, are permitted provided that the following conditions are met:
23
+
24
+ * Redistributions of source code must retain the above copyright notice, this
25
+ list of conditions and the following disclaimer.
26
+
27
+ * Redistributions in binary form must reproduce the above copyright notice,
28
+ this list of conditions and the following disclaimer in the documentation
29
+ and/or other materials provided with the distribution.
30
+
31
+ * Neither the name of the copyright holder nor the names of its
32
+ contributors may be used to endorse or promote products derived from
33
+ this software without specific prior written permission.
34
+
35
+ THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
36
+ AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
37
+ IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
38
+ DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
39
+ FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
40
+ DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
41
+ SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
42
+ CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
43
+ OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
44
+ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
45
+ ```
46
+
47
+ ## pix2pix — pytorch-CycleGAN-and-pix2pix
48
+
49
+ `tests/models/pix2pix.py` and `examples/pix2pix.ipynb` contain the U-Net generator and the
50
+ PatchGAN discriminator as defined in `models/networks.py` of the authors' implementation
51
+ (<https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix>).
52
+
53
+ ```text
54
+ Copyright (c) 2017, Jun-Yan Zhu and Taesung Park
55
+ All rights reserved.
56
+
57
+ Redistribution and use in source and binary forms, with or without
58
+ modification, are permitted provided that the following conditions are met:
59
+
60
+ * Redistributions of source code must retain the above copyright notice, this
61
+ list of conditions and the following disclaimer.
62
+
63
+ * Redistributions in binary form must reproduce the above copyright notice,
64
+ this list of conditions and the following disclaimer in the documentation
65
+ and/or other materials provided with the distribution.
66
+
67
+ THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
68
+ AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
69
+ IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
70
+ DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
71
+ FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
72
+ DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
73
+ SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
74
+ CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
75
+ OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
76
+ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
77
+
78
+
79
+ --------------------------- LICENSE FOR pix2pix --------------------------------
80
+ BSD License
81
+
82
+ For pix2pix software
83
+ Copyright (c) 2016, Phillip Isola and Jun-Yan Zhu
84
+ All rights reserved.
85
+
86
+ Redistribution and use in source and binary forms, with or without
87
+ modification, are permitted provided that the following conditions are met:
88
+
89
+ * Redistributions of source code must retain the above copyright notice, this
90
+ list of conditions and the following disclaimer.
91
+
92
+ * Redistributions in binary form must reproduce the above copyright notice,
93
+ this list of conditions and the following disclaimer in the documentation
94
+ and/or other materials provided with the distribution.
95
+
96
+ ----------------------------- LICENSE FOR DCGAN --------------------------------
97
+ BSD License
98
+
99
+ For dcgan.torch software
100
+
101
+ Copyright (c) 2015, Facebook, Inc. All rights reserved.
102
+
103
+ Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
104
+
105
+ Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
106
+
107
+ Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
108
+
109
+ Neither the name Facebook nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
110
+
111
+ THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
112
+ ```
@@ -0,0 +1,117 @@
1
+ [project]
2
+ name = "tensorgami"
3
+ dynamic = ["version", "readme"] # version: src/tensorgami/__init__.py; readme: README.md, below
4
+ description = "Draw the architecture and information flow of PyTorch models: figures for papers, reports and slides"
5
+ license = "MIT"
6
+ license-files = ["LICENSE"]
7
+ authors = [{ name = "Andrei Utina" }]
8
+ keywords = [
9
+ "pytorch", "neural network", "architecture", "diagram", "visualization", "deep learning",
10
+ "figure", "matplotlib",
11
+ ]
12
+ classifiers = [
13
+ "Development Status :: 3 - Alpha",
14
+ "Framework :: Jupyter",
15
+ "Intended Audience :: Science/Research",
16
+ "Intended Audience :: Education",
17
+ "Operating System :: OS Independent",
18
+ "Programming Language :: Python :: 3",
19
+ "Programming Language :: Python :: 3 :: Only",
20
+ "Programming Language :: Python :: 3.11",
21
+ "Programming Language :: Python :: 3.12",
22
+ "Programming Language :: Python :: 3.13",
23
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
24
+ "Topic :: Scientific/Engineering :: Visualization",
25
+ "Topic :: Multimedia :: Graphics :: Presentation",
26
+ ]
27
+ requires-python = ">=3.11"
28
+ dependencies = [
29
+ "matplotlib>=3.7", # figures, PNG, Jupyter display (ADR-0005); the only runtime dependency
30
+ ]
31
+
32
+ [project.optional-dependencies]
33
+ torch = ["torch>=2.2"] # never a hard dependency: users keep their own CPU/CUDA build (ADR-0006)
34
+
35
+ [project.urls]
36
+ Homepage = "https://github.com/AndreiUtina/TensorGami"
37
+ Repository = "https://github.com/AndreiUtina/TensorGami"
38
+ Issues = "https://github.com/AndreiUtina/TensorGami/issues"
39
+ Changelog = "https://github.com/AndreiUtina/TensorGami/blob/main/CHANGELOG.md"
40
+ Gallery = "https://github.com/AndreiUtina/TensorGami/blob/main/docs/gallery/README.md"
41
+
42
+ [dependency-groups]
43
+ dev = [
44
+ "pytest>=8",
45
+ "ruff>=0.6",
46
+ ]
47
+ docs = [
48
+ "typst>=0.12",
49
+ ]
50
+ torch = [
51
+ # the newest torch with Intel-Mac wheels is 2.2.2, which needs numpy < 2 (ADR-0006)
52
+ "torch==2.2.2; sys_platform == 'darwin' and platform_machine == 'x86_64'",
53
+ "numpy<2; sys_platform == 'darwin' and platform_machine == 'x86_64'",
54
+ "torch>=2.2; sys_platform != 'darwin' or platform_machine != 'x86_64'",
55
+ ]
56
+
57
+ [build-system]
58
+ requires = ["hatchling", "hatch-fancy-pypi-readme>=24"]
59
+ build-backend = "hatchling.build"
60
+
61
+ [tool.hatch.version]
62
+ path = "src/tensorgami/__init__.py"
63
+
64
+ # README.md links to the repository's files by relative paths, which GitHub resolves in a
65
+ # private repository and on every branch; on PyPI they become absolute (images: the raw files)
66
+ [tool.hatch.metadata.hooks.fancy-pypi-readme]
67
+ content-type = "text/markdown"
68
+
69
+ [[tool.hatch.metadata.hooks.fancy-pypi-readme.fragments]]
70
+ path = "README.md"
71
+
72
+ [[tool.hatch.metadata.hooks.fancy-pypi-readme.substitutions]]
73
+ pattern = '(!\[[^\]]*\]\()(?!https?://)([^)\s]+)\)'
74
+ replacement = '\1https://raw.githubusercontent.com/AndreiUtina/TensorGami/main/\2)'
75
+
76
+ [[tool.hatch.metadata.hooks.fancy-pypi-readme.substitutions]]
77
+ pattern = '(\]\()(?!https?://|#)([^)\s]+)\)'
78
+ replacement = '\1https://github.com/AndreiUtina/TensorGami/blob/main/\2)'
79
+
80
+ [tool.hatch.build.targets.wheel]
81
+ packages = ["src/tensorgami"]
82
+
83
+ [tool.hatch.build.targets.sdist] # the package and what describes it; docs and examples stay on GitHub
84
+ include = [
85
+ "/src/tensorgami",
86
+ "/README.md",
87
+ "/LICENSE",
88
+ "/THIRD_PARTY_NOTICES.md",
89
+ "/CHANGELOG.md",
90
+ "/CITATION.cff",
91
+ "/pyproject.toml",
92
+ ]
93
+
94
+ [tool.uv]
95
+ default-groups = ["dev", "docs", "torch"]
96
+ environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'", "sys_platform == 'win32'"]
97
+
98
+ [tool.uv.sources]
99
+ torch = [{ index = "pytorch-cpu", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }] # no CUDA in dev/CI
100
+
101
+ [[tool.uv.index]]
102
+ name = "pytorch-cpu"
103
+ url = "https://download.pytorch.org/whl/cpu"
104
+ explicit = true
105
+
106
+ [tool.pytest.ini_options]
107
+ testpaths = ["tests"]
108
+
109
+ [tool.ruff]
110
+ line-length = 100
111
+ target-version = "py311"
112
+
113
+ [tool.ruff.lint]
114
+ select = ["E", "F", "I", "B", "UP"]
115
+
116
+ [tool.ruff.format]
117
+ exclude = ["*.md", "docs/survey/*"] # prose code blocks and the survey's build scripts stay as written