visualtorch 1.1.0__tar.gz → 1.2.0__tar.gz

{visualtorch-1.1.0/visualtorch.egg-info → visualtorch-1.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: visualtorch
-Version: 1.1.0
+Version: 1.2.0
 Summary: Architecture visualization of Torch models
 Home-page: https://github.com/willyfh/visualtorch
 Author: Willy Fitra Hendria
@@ -54,21 +54,22 @@ Dynamic: summary
 
 </div>
 
-**VisualTorch** aims to help visualize Torch-based neural network architectures. It currently supports generating flow-style, graph-style, and LeNet-style architectures for PyTorch Sequential and Custom models. This tool is adapted from [visualkeras](https://github.com/paulgavrikov/visualkeras), [pytorchviz](https://github.com/szagoruyko/pytorchviz), and [pytorch-summary](https://github.com/sksq96/pytorch-summary).
+**VisualTorch** aims to help visualize Torch-based neural network architectures. It currently supports generating flow-style, graph-style, and LeNet-style architectures for PyTorch Sequential and Custom models. Its original visual styles were inspired by [visualkeras](https://github.com/paulgavrikov/visualkeras), [pytorchviz](https://github.com/szagoruyko/pytorchviz), [pytorch-summary](https://github.com/sksq96/pytorch-summary), and [torchview](https://github.com/mert-kurttutan/torchview); since then, it has grown its own unified tracing backend and architecture-handling logic well beyond its origins.
 
 **Note:** `1.0+` is a major release with breaking API changes, but with significantly better features and algorithms - upgrading is recommended. For the old API, use `0.2.5` or older.
 
-**Limitation:** VisualTorch traces a real forward pass to build the diagram, which has two inherent
-limitations shared by any tracing-based approach (not bugs, and not fixable without full symbolic
-execution): (1) models with **data-dependent control flow** (e.g. a branch only taken if a tensor
-value crosses some threshold) only show whichever branch the traced dummy input happened to take;
-(2) a layer that returns **multiple meaningful output tensors** (e.g. a custom multi-task head)
-only has its first tensor's shape reflected in that node's size/label - its downstream connections
-are still correct either way. Contributions are welcome!
+**Limitation:** VisualTorch traces a real forward pass to build the diagram, which has an inherent
+limitation shared by any tracing-based approach (not a bug, and not fixable without full symbolic
+execution): models with **data-dependent control flow** (e.g. a branch only taken if a tensor
+value crosses some threshold) only show whichever branch the traced dummy input happened to take.
+Separately, a layer that returns **multiple meaningful output tensors** (e.g. a custom multi-task
+head, or `nn.LSTM`'s `(output, (h_n, c_n))`) still has its node's size based on only its first
+tensor; with `show_dimension=True`, every output tensor's shape is shown in the label, not just
+the first. Downstream connections are correct either way. Contributions are welcome!
 
 <div align="center">
 
-![VisualTorch Examples](docs/source/_static/images/banners/readme-examples.png)
+![VisualTorch Examples](https://raw.githubusercontent.com/willyfh/visualtorch/e6ad79751e0f7412b1074beb45f9baeccd1419e4/docs/source/_static/images/banners/readme-examples.png)
 
 </div>
 
@@ -100,16 +101,16 @@ Please feel free to send a pull request to contribute to this project by followi
 
 This poject is available as open source under the terms of the [MIT License](https://github.com/willyfh/visualtorch/blob/main/LICENSE).
 
-Originally, this project was based on the [visualkeras](https://github.com/paulgavrikov/visualkeras) (under the MIT license), with additional modifications inspired by [pytorchviz](https://github.com/szagoruyko/pytorchviz), and [pytorch-summary](https://github.com/sksq96/pytorch-summary), both of which are also licensed under the MIT license.
+Originally, this project was based on the [visualkeras](https://github.com/paulgavrikov/visualkeras) (under the MIT license), with additional modifications inspired by [pytorchviz](https://github.com/szagoruyko/pytorchviz), [pytorch-summary](https://github.com/sksq96/pytorch-summary), and [torchview](https://github.com/mert-kurttutan/torchview), all of which are also licensed under the MIT license.
 
 ## Citation
 
 Please cite this project in your publications if it helps your research.
 
-**Note:** the paper below describes the API as of its publication date (2024). VisualTorch has
-since had breaking API changes (see the [documentation](https://visualtorch.readthedocs.io/en/latest/)
-for the current API) - the DOI always resolves to what was actually reviewed and published, so
-it isn't updated to match.
+**Note:** the paper below describes VisualTorch as of its publication date (2024). The project has
+since been substantially refactored, including breaking API changes (see the
+[documentation](https://visualtorch.readthedocs.io/en/latest/) for the current API) - the DOI
+always resolves to what was actually reviewed and published.
 
 ```bibtex
 @article{Hendria2024,

{visualtorch-1.1.0 → visualtorch-1.2.0}/README.md

@@ -5,21 +5,22 @@
 
 </div>
 
-**VisualTorch** aims to help visualize Torch-based neural network architectures. It currently supports generating flow-style, graph-style, and LeNet-style architectures for PyTorch Sequential and Custom models. This tool is adapted from [visualkeras](https://github.com/paulgavrikov/visualkeras), [pytorchviz](https://github.com/szagoruyko/pytorchviz), and [pytorch-summary](https://github.com/sksq96/pytorch-summary).
+**VisualTorch** aims to help visualize Torch-based neural network architectures. It currently supports generating flow-style, graph-style, and LeNet-style architectures for PyTorch Sequential and Custom models. Its original visual styles were inspired by [visualkeras](https://github.com/paulgavrikov/visualkeras), [pytorchviz](https://github.com/szagoruyko/pytorchviz), [pytorch-summary](https://github.com/sksq96/pytorch-summary), and [torchview](https://github.com/mert-kurttutan/torchview); since then, it has grown its own unified tracing backend and architecture-handling logic well beyond its origins.
 
 **Note:** `1.0+` is a major release with breaking API changes, but with significantly better features and algorithms - upgrading is recommended. For the old API, use `0.2.5` or older.
 
-**Limitation:** VisualTorch traces a real forward pass to build the diagram, which has two inherent
-limitations shared by any tracing-based approach (not bugs, and not fixable without full symbolic
-execution): (1) models with **data-dependent control flow** (e.g. a branch only taken if a tensor
-value crosses some threshold) only show whichever branch the traced dummy input happened to take;
-(2) a layer that returns **multiple meaningful output tensors** (e.g. a custom multi-task head)
-only has its first tensor's shape reflected in that node's size/label - its downstream connections
-are still correct either way. Contributions are welcome!
+**Limitation:** VisualTorch traces a real forward pass to build the diagram, which has an inherent
+limitation shared by any tracing-based approach (not a bug, and not fixable without full symbolic
+execution): models with **data-dependent control flow** (e.g. a branch only taken if a tensor
+value crosses some threshold) only show whichever branch the traced dummy input happened to take.
+Separately, a layer that returns **multiple meaningful output tensors** (e.g. a custom multi-task
+head, or `nn.LSTM`'s `(output, (h_n, c_n))`) still has its node's size based on only its first
+tensor; with `show_dimension=True`, every output tensor's shape is shown in the label, not just
+the first. Downstream connections are correct either way. Contributions are welcome!
 
 <div align="center">
 
-![VisualTorch Examples](docs/source/_static/images/banners/readme-examples.png)
+![VisualTorch Examples](https://raw.githubusercontent.com/willyfh/visualtorch/e6ad79751e0f7412b1074beb45f9baeccd1419e4/docs/source/_static/images/banners/readme-examples.png)
 
 </div>
 
@@ -51,16 +52,16 @@ Please feel free to send a pull request to contribute to this project by followi
 
 This poject is available as open source under the terms of the [MIT License](https://github.com/willyfh/visualtorch/blob/main/LICENSE).
 
-Originally, this project was based on the [visualkeras](https://github.com/paulgavrikov/visualkeras) (under the MIT license), with additional modifications inspired by [pytorchviz](https://github.com/szagoruyko/pytorchviz), and [pytorch-summary](https://github.com/sksq96/pytorch-summary), both of which are also licensed under the MIT license.
+Originally, this project was based on the [visualkeras](https://github.com/paulgavrikov/visualkeras) (under the MIT license), with additional modifications inspired by [pytorchviz](https://github.com/szagoruyko/pytorchviz), [pytorch-summary](https://github.com/sksq96/pytorch-summary), and [torchview](https://github.com/mert-kurttutan/torchview), all of which are also licensed under the MIT license.
 
 ## Citation
 
 Please cite this project in your publications if it helps your research.
 
-**Note:** the paper below describes the API as of its publication date (2024). VisualTorch has
-since had breaking API changes (see the [documentation](https://visualtorch.readthedocs.io/en/latest/)
-for the current API) - the DOI always resolves to what was actually reviewed and published, so
-it isn't updated to match.
+**Note:** the paper below describes VisualTorch as of its publication date (2024). The project has
+since been substantially refactored, including breaking API changes (see the
+[documentation](https://visualtorch.readthedocs.io/en/latest/) for the current API) - the DOI
+always resolves to what was actually reviewed and published.
 
 ```bibtex
 @article{Hendria2024,

{visualtorch-1.1.0 → visualtorch-1.2.0}/setup.py

@@ -21,7 +21,7 @@ def _read_requirements(file: str) -> list:
 
 setuptools.setup(
     name="visualtorch",
-    version="1.1.0",
+    version="1.2.0",
     author="Willy Fitra Hendria",
     author_email="willyfitrahendria@gmail.com",
     description="Architecture visualization of Torch models",

{visualtorch-1.1.0 → visualtorch-1.2.0}/tests/test_flow.py

@@ -174,16 +174,16 @@ def test_rnn_model_flow_view_runs(rnn_model: nn.Module) -> None:
 
 
 @pytest.mark.parametrize("orientation", ["x", "y", "z"])
-def test_flow_view_one_dim_orientation(classifier_model: nn.Module, orientation: str) -> None:
+def test_flow_view_low_dim_orientation(classifier_model: nn.Module, orientation: str) -> None:
     """Test flow view on a model with a 1D output, for every supported orientation."""
-    img = flow_view(classifier_model, input_shape=(1, 3, 16, 16), one_dim_orientation=orientation)
+    img = flow_view(classifier_model, input_shape=(1, 3, 16, 16), low_dim_orientation=orientation)
     assert img is not None
 
 
-def test_flow_view_invalid_one_dim_orientation_raises(classifier_model: nn.Module) -> None:
-    """An unsupported one_dim_orientation should raise a clear ValueError."""
+def test_flow_view_invalid_low_dim_orientation_raises(classifier_model: nn.Module) -> None:
+    """An unsupported low_dim_orientation should raise a clear ValueError."""
     with pytest.raises(ValueError, match="unsupported orientation"):
-        flow_view(classifier_model, input_shape=(1, 3, 16, 16), one_dim_orientation="bad")
+        flow_view(classifier_model, input_shape=(1, 3, 16, 16), low_dim_orientation="bad")
 
 
 def test_flow_view_with_type_ignore(sequential_model: nn.Sequential) -> None:
@@ -535,3 +535,54 @@ def test_flow_view_mismatched_depth_siamese_branches_needs_no_detour() -> None:
     img_matched = flow_view(SiameseNetDepthMatched(), input_shape=input_shape)
 
     assert img_mismatched.size[1] == img_matched.size[1]
+
+
+def test_flow_view_low_dim_orientation_affects_2d_shapes() -> None:
+    """A 2D shape (e.g. an RNN's (seq_len, hidden_size)) should now respond to
+    low_dim_orientation too, not just genuine 1D shapes.
+    """  # noqa: D205
+
+    class SequenceClassifier(nn.Module):
+        def __init__(self, hidden_size: int) -> None:
+            super().__init__()
+            self.lstm = nn.LSTM(input_size=8, hidden_size=hidden_size, batch_first=True)
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            out, _ = self.lstm(x)
+            return out
+
+    model = SequenceClassifier(hidden_size=64)
+    input_shape = (1, 5, 8)
+
+    sizes = {
+        orientation: flow_view(model, input_shape=input_shape, low_dim_orientation=orientation).size
+        for orientation in ("x", "y", "z")
+    }
+
+    assert len(set(sizes.values())) == 3, f"expected all 3 orientations to differ, got {sizes}"
+
+
+def test_flow_view_2d_shape_seq_len_is_discarded() -> None:
+    """The positional-like dim (e.g. seq_len) of a 2D shape shouldn't affect box size -
+    only the feature-like dim (e.g. hidden_size) should.
+    """  # noqa: D205
+
+    class SequenceClassifier(nn.Module):
+        def __init__(self, hidden_size: int) -> None:
+            super().__init__()
+            self.lstm = nn.LSTM(input_size=8, hidden_size=hidden_size, batch_first=True)
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            out, _ = self.lstm(x)
+            return out
+
+    model = SequenceClassifier(hidden_size=64)
+    img_short_seq = flow_view(model, input_shape=(1, 5, 8))
+    img_long_seq = flow_view(model, input_shape=(1, 50, 8))
+
+    assert img_short_seq.tobytes() == img_long_seq.tobytes()
+
+    model_bigger_hidden = SequenceClassifier(hidden_size=256)
+    img_bigger_hidden = flow_view(model_bigger_hidden, input_shape=(1, 5, 8))
+
+    assert img_short_seq.tobytes() != img_bigger_hidden.tobytes()

{visualtorch-1.1.0 → visualtorch-1.2.0}/tests/test_lenet_style.py

@@ -174,16 +174,16 @@ def test_rnn_model_lenet_view_runs(rnn_model: nn.Module) -> None:
 
 
 @pytest.mark.parametrize("orientation", ["x", "y", "z"])
-def test_lenet_view_one_dim_orientation(classifier_model: nn.Module, orientation: str) -> None:
+def test_lenet_view_low_dim_orientation(classifier_model: nn.Module, orientation: str) -> None:
     """Test lenet view on a model with a 1D output, for every supported orientation."""
-    img = lenet_view(classifier_model, input_shape=(1, 3, 16, 16), one_dim_orientation=orientation)
+    img = lenet_view(classifier_model, input_shape=(1, 3, 16, 16), low_dim_orientation=orientation)
     assert img is not None
 
 
-def test_lenet_view_invalid_one_dim_orientation_raises(classifier_model: nn.Module) -> None:
-    """An unsupported one_dim_orientation should raise a clear ValueError."""
+def test_lenet_view_invalid_low_dim_orientation_raises(classifier_model: nn.Module) -> None:
+    """An unsupported low_dim_orientation should raise a clear ValueError."""
     with pytest.raises(ValueError, match="unsupported orientation"):
-        lenet_view(classifier_model, input_shape=(1, 3, 16, 16), one_dim_orientation="bad")
+        lenet_view(classifier_model, input_shape=(1, 3, 16, 16), low_dim_orientation="bad")
 
 
 def test_lenet_view_with_type_ignore(sequential_model: nn.Sequential) -> None:
@@ -328,3 +328,54 @@ def test_lenet_view_funnels_survive_large_de_differences_between_layers() -> Non
     non_bg = int((np.array(img.convert("RGB")) != 255).any(axis=2).sum())
     error_msg = f"non-background pixel count {non_bg} outside expected range - funnel likely broken"
     assert 110000 <= non_bg <= 145000, error_msg
+
+
+def test_lenet_view_low_dim_orientation_affects_2d_shapes() -> None:
+    """A 2D shape (e.g. an RNN's (seq_len, hidden_size)) should now respond to
+    low_dim_orientation too, not just genuine 1D shapes.
+    """  # noqa: D205
+
+    class SequenceClassifier(nn.Module):
+        def __init__(self, hidden_size: int) -> None:
+            super().__init__()
+            self.lstm = nn.LSTM(input_size=8, hidden_size=hidden_size, batch_first=True)
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            out, _ = self.lstm(x)
+            return out
+
+    model = SequenceClassifier(hidden_size=64)
+    input_shape = (1, 5, 8)
+
+    sizes = {
+        orientation: lenet_view(model, input_shape=input_shape, low_dim_orientation=orientation).size
+        for orientation in ("x", "y", "z")
+    }
+
+    assert len(set(sizes.values())) == 3, f"expected all 3 orientations to differ, got {sizes}"
+
+
+def test_lenet_view_2d_shape_seq_len_is_discarded() -> None:
+    """The positional-like dim (e.g. seq_len) of a 2D shape shouldn't affect box size -
+    only the feature-like dim (e.g. hidden_size) should.
+    """  # noqa: D205
+
+    class SequenceClassifier(nn.Module):
+        def __init__(self, hidden_size: int) -> None:
+            super().__init__()
+            self.lstm = nn.LSTM(input_size=8, hidden_size=hidden_size, batch_first=True)
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            out, _ = self.lstm(x)
+            return out
+
+    model = SequenceClassifier(hidden_size=64)
+    img_short_seq = lenet_view(model, input_shape=(1, 5, 8))
+    img_long_seq = lenet_view(model, input_shape=(1, 50, 8))
+
+    assert img_short_seq.tobytes() == img_long_seq.tobytes()
+
+    model_bigger_hidden = SequenceClassifier(hidden_size=256)
+    img_bigger_hidden = lenet_view(model_bigger_hidden, input_shape=(1, 5, 8))
+
+    assert img_short_seq.tobytes() != img_bigger_hidden.tobytes()

{visualtorch-1.1.0 → visualtorch-1.2.0}/tests/test_regression_issues.py

@@ -11,8 +11,9 @@ import torch
 from torch import nn
 from visualtorch.backend import extract_architecture
 from visualtorch.flow import flow_view
+from visualtorch.graph import graph_view
 from visualtorch.lenet_style import lenet_view
-from visualtorch.utils.utils import self_multiply
+from visualtorch.utils.utils import format_shape_label, self_multiply
 
 
 @pytest.fixture()
@@ -112,3 +113,48 @@ def test_flow_view_recurrent_sequence_length_does_not_inflate_diagram_height(rec
     long_img = flow_view(model, input_shape=(1, 200, 10))
 
     assert long_img.height == short_img.height
+
+
+@pytest.fixture()
+def lstm_using_hidden_state_model() -> nn.Module:
+    """A model that consumes `nn.LSTM`'s hidden state (h_n), not its sequence output.
+
+    `nn.LSTM.forward()` returns `(output, (h_n, c_n))` - three tensors, not one. Before the fix,
+    only `output`'s shape was ever recorded, even when (as here) the model actually uses `h_n`
+    instead - silently dropping the shape of the tensor that matters, with nothing shown to
+    indicate more than one tensor even exists.
+    """
+
+    class Model(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.lstm = nn.LSTM(input_size=10, hidden_size=20, batch_first=True)
+            self.fc = nn.Linear(20, 5)
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            _output, (h_n, _c_n) = self.lstm(x)
+            return self.fc(h_n.squeeze(0))
+
+    return Model()
+
+
+def test_extract_architecture_records_every_output_tensor_shape(lstm_using_hidden_state_model: nn.Module) -> None:
+    """A multi-output leaf module's TracedLayer should record all three of its output shapes."""
+    architecture = extract_architecture(lstm_using_hidden_state_model, (1, 7, 10))
+    lstm_layer = next(layer for column in architecture.columns for layer in column if isinstance(layer.module, nn.LSTM))
+
+    assert lstm_layer.output_shape == (1, 7, 20)
+    assert lstm_layer.extra_output_shapes == ((1, 1, 20), (1, 1, 20))
+
+
+def test_format_shape_label_appends_extra_shapes() -> None:
+    """format_shape_label should append every extra shape, and omit the `+` entirely when there are none."""
+    assert format_shape_label((1, 7, 20), ()) == "(1, 7, 20)"
+    assert format_shape_label((1, 7, 20), ((1, 1, 20), (1, 1, 20))) == "(1, 7, 20) + (1, 1, 20) + (1, 1, 20)"
+
+
+@pytest.mark.parametrize("view", [graph_view, flow_view, lenet_view])
+def test_show_dimension_includes_every_output_shape(lstm_using_hidden_state_model: nn.Module, view: object) -> None:
+    """show_dimension=True shouldn't crash, and should still work, for a multi-output leaf module."""
+    img = view(lstm_using_hidden_state_model, input_shape=(1, 7, 10), show_dimension=True)  # type: ignore[operator]
+    assert img is not None

{visualtorch-1.1.0 → visualtorch-1.2.0}/tests/test_render.py

@@ -3,11 +3,15 @@
 # Copyright (C) 2024 Willy Fitra Hendria
 # SPDX-License-Identifier: MIT
 
+from collections import defaultdict
+
 import pytest
 import torch
 from torch import nn
 from visualtorch import render
 from visualtorch.backend import extract_architecture
+from visualtorch.utils.layer_utils import Input
+from visualtorch.utils.utils import PALETTES
 
 
 @pytest.fixture()
@@ -148,3 +152,43 @@ def test_render_handles_unused_input_tensor() -> None:
 
     img = render(PartiallyUnusedNet(), input_shape=((1, 10), (1, 5)), style="graph")
     assert img is not None
+
+
+@pytest.mark.parametrize("palette", sorted(PALETTES))
+def test_render_runs_for_every_named_palette(sequential_model: nn.Sequential, palette: str) -> None:
+    """Every named palette should render without error - catches any malformed hex color."""
+    img = render(sequential_model, input_shape=(1, 3, 16, 16), style="graph", palette=palette)
+    assert img is not None
+
+
+def test_render_rejects_unsupported_palette(sequential_model: nn.Sequential) -> None:
+    """An unrecognized palette name should raise a clear error, not silently fall back."""
+    with pytest.raises(ValueError, match="Unsupported palette"):
+        render(sequential_model, input_shape=(1, 3, 16, 16), style="graph", palette="bogus")
+
+
+def test_render_palette_changes_fallback_colors(sequential_model: nn.Sequential) -> None:
+    """A different palette should actually change the colors of unmapped layer types."""
+    default = render(sequential_model, input_shape=(1, 3, 16, 16), style="graph")
+    dracula = render(sequential_model, input_shape=(1, 3, 16, 16), style="graph", palette="dracula")
+
+    assert default.tobytes() != dracula.tobytes()
+
+
+def test_render_color_map_overrides_palette(sequential_model: nn.Sequential) -> None:
+    """An explicit color_map entry should still win over the palette fallback."""
+    color_map: dict = defaultdict(dict)
+    color_map[Input]["fill"] = "#abcdef"
+    color_map[nn.Conv2d]["fill"] = "#123456"
+    color_map[nn.ReLU]["fill"] = "#654321"
+
+    okabe_ito = render(sequential_model, input_shape=(1, 3, 16, 16), style="graph", color_map=color_map)
+    dracula = render(
+        sequential_model,
+        input_shape=(1, 3, 16, 16),
+        style="graph",
+        color_map=color_map,
+        palette="dracula",
+    )
+
+    assert okabe_ito.tobytes() == dracula.tobytes()

{visualtorch-1.1.0 → visualtorch-1.2.0}/visualtorch/__init__.py

@@ -11,6 +11,7 @@ from visualtorch.render import (
     render,
 )
 from visualtorch.utils.layer_utils import Input
+from visualtorch.utils.utils import PALETTES
 
 __all__ = [
     "render",
@@ -19,4 +20,5 @@ __all__ = [
     "FlowStyleOptions",
     "LenetStyleOptions",
     "Input",
+    "PALETTES",
 ]

{visualtorch-1.1.0 → visualtorch-1.2.0}/visualtorch/backend.py

@@ -67,7 +67,10 @@ def extract_architecture(model: nn.Module, input_shape: InputShape) -> Architect
     """
     input_shapes = validate_input_shape(input_shape)
 
-    id_to_module, id_to_output_shape, edges, input_ids = trace_module_graph(model, input_shapes)
+    id_to_module, id_to_output_shape, id_to_extra_output_shapes, edges, input_ids = trace_module_graph(
+        model,
+        input_shapes,
+    )
 
     nodes = list(id_to_module.keys())
     id_to_index = {node_id: idx for idx, node_id in enumerate(nodes)}
@@ -114,6 +117,7 @@ def extract_architecture(model: nn.Module, input_shape: InputShape) -> Architect
             module=id_to_module[node_id],
             output_shape=id_to_output_shape[node_id],
             node_id=node_id,
+            extra_output_shapes=id_to_extra_output_shapes.get(node_id, ()),
         )
         columns[depth[node_id]].append(wrapper)
 

{visualtorch-1.1.0 → visualtorch-1.2.0}/visualtorch/flow.py

@@ -22,8 +22,10 @@ from .utils.utils import (
     ColorWheel,
     ImageDraw,
     InputShape,
+    format_shape_label,
     get_rgba_tuple,
     linear_layout,
+    resolve_palette,
     self_multiply,
     vertical_image_concat,
 )
@@ -41,7 +43,8 @@ def flow_view(
     scale_xy: float = 1,
     type_ignore: list | None = None,
     color_map: dict | None = None,
-    one_dim_orientation: str = "z",
+    palette: str = "okabe_ito",
+    low_dim_orientation: str = "z",
     background_fill: str | tuple[int, ...] = "white",
     draw_volume: bool = True,
     padding: int = 10,
@@ -75,7 +78,13 @@ def flow_view(
         type_ignore (list, optional): List of layer types in the torch model to ignore during drawing.
         color_map (dict, optional): Dictionary defining fill and outline colors for each layer by class type.
             Will fallback to default values for unspecified classes.
-        one_dim_orientation (str, optional): Axis on which one-dim layers should be drawn. E.g., 'x', 'y', or 'z'.
+        palette (str, optional): Named color palette used as the fallback for any layer type not
+            given an explicit override via `color_map`. One of `"okabe_ito"` (default,
+            colorblind-safe), `"tol_bright"`, `"tol_muted"`, `"tab10"`, `"grayscale"`, `"nord"`,
+            `"dracula"`, `"gruvbox"`, `"solarized"`, `"material"`, `"catppuccin"`.
+        low_dim_orientation (str, optional): Axis on which a layer without real spatial/channel
+            structure (a 1D shape, or a 2D shape like an RNN/attention layer's
+            `(seq_len, hidden_size)`) should be drawn. One of `'x'`, `'y'`, or `'z'`.
         background_fill (str or tuple, optional): Background color for the image. A string or a tuple (R, G, B, A).
         draw_volume (bool, optional): Flag to switch between 3D volumetric view and 2D box view.
         padding (int, optional): Distance in pixels before the first and after the last layer.
@@ -132,9 +141,9 @@ def flow_view(
     filtered_columns = [column for column in filtered_columns if column]
 
     layer_types: list[type] = []
-    color_wheel = ColorWheel()
+    color_wheel = ColorWheel(colors=resolve_palette(palette))
     make_box = _box_factory(
-        one_dim_orientation,
+        low_dim_orientation,
         scale_xy,
         min_xy,
         max_xy,
@@ -232,7 +241,7 @@ def flow_view(
 
 
 def _box_factory(
-    one_dim_orientation: str,
+    low_dim_orientation: str,
     scale_xy: float,
     min_xy: int,
     max_xy: int,
@@ -251,19 +260,20 @@ def _box_factory(
     def make_box(layer: TracedLayer) -> Box:
         shape = layer.output_shape[1:]  # drop batch size
 
-        if len(shape) == 1:
-            if one_dim_orientation in ("x", "y", "z"):
-                shape = (1,) * "cxyz".index(one_dim_orientation) + shape
+        if len(shape) in (1, 2):
+            # Neither a 1D nor a 2D shape has real spatial/channel structure - there's nothing
+            # to distinguish "channel" from "spatial" the way a genuine (C, H, W) feature map
+            # does. Take the last value (for 2D, e.g. an RNN/attention layer's
+            # (seq_len, hidden_size), this is the feature/channel-like one, matching PyTorch's
+            # (..., seq, feature) convention for sequence data; for 1D it's the only value) and
+            # let the user place it on whichever axis they choose, same as any 1D value - the
+            # positional-like dim, if any, is discarded either way.
+            value = shape[-1]
+            if low_dim_orientation in ("x", "y", "z"):
+                shape = (1,) * "cxyz".index(low_dim_orientation) + (value,)
             else:
-                error_msg = f"unsupported orientation: {one_dim_orientation}"
+                error_msg = f"unsupported orientation: {low_dim_orientation}"
                 raise ValueError(error_msg)
-        elif len(shape) == 2:
-            # A 2D non-batch shape (e.g. (seq_len, hidden_size) from an RNN/attention layer)
-            # isn't a CNN feature map missing a channel dim - there's no channel axis at all.
-            # Box's "3D" skew (de, below) is driven by shape[1], so a dummy 1 goes there
-            # instead of either real dim, keeping the two real dims on the box's actual width
-            # and height instead of one of them inflating the skew for a long sequence.
-            shape = (shape[0], 1, shape[1])
 
         ori_shape = shape
         shape = shape + (1,) * (4 - len(shape))  # expand 4D.
@@ -278,6 +288,7 @@ def _box_factory(
 
         box = Box()
         box.output_shape = tuple(ori_shape)
+        box.extra_output_shapes = layer.extra_output_shapes
         box.de = int(x / 3) if draw_volume else 0
 
         box.x1 = 0
@@ -476,7 +487,7 @@ def _draw_legend(
 
 def _column_label_and_center(column: list[VolumetricBox]) -> tuple[str, float]:
     """A column's shape label (joined across branches) and its shared x-center."""
-    label = " / ".join(str(box.output_shape) for box in column)
+    label = " / ".join(format_shape_label(box.output_shape, box.extra_output_shapes) for box in column)
     center_x = (column[0].x1 + column[0].x2) / 2
     return label, center_x
 

{visualtorch-1.1.0 → visualtorch-1.2.0}/visualtorch/graph.py

@@ -15,7 +15,7 @@ from PIL import Image, ImageFont
 from .backend import extract_architecture
 from .connectors import compute_skip_levels, draw_connector
 from .utils.traced_layer import TracedLayer
-from .utils.utils import Box, Circle, ColorWheel, Ellipses, ImageDraw, InputShape
+from .utils.utils import Box, Circle, ColorWheel, Ellipses, ImageDraw, InputShape, format_shape_label, resolve_palette
 
 
 def graph_view(
@@ -23,6 +23,7 @@ def graph_view(
     input_shape: InputShape,
     to_file: str | None = None,
     color_map: dict[Any, Any] | None = None,
+    palette: str = "okabe_ito",
     node_size: int = 50,
     background_fill: str | tuple[int, ...] = "white",
     padding: int = 10,
@@ -51,6 +52,10 @@ def graph_view(
             will disable writing.
         color_map (dict, optional): Dict defining fill and outline for each layer by class type. Will fallback
             to default values for not specified classes.
+        palette (str, optional): Named color palette used as the fallback for any layer type not
+            given an explicit override via `color_map`. One of `"okabe_ito"` (default,
+            colorblind-safe), `"tol_bright"`, `"tol_muted"`, `"tab10"`, `"grayscale"`, `"nord"`,
+            `"dracula"`, `"gruvbox"`, `"solarized"`, `"material"`, `"catppuccin"`.
         node_size (int, optional): Size in pixels each node will have.
         background_fill (Any, optional): Color for the image background. Can be str or (R,G,B,A).
         padding (int, optional): Distance in pixels before the first and after the last layer.
@@ -111,7 +116,7 @@ def graph_view(
         _color_map,
         opacity,
         layer_spacing,
-        ColorWheel(),
+        ColorWheel(colors=resolve_palette(palette)),
     )
 
     # An edge whose endpoint was just dropped above (a hidden input's own edges) can no longer
@@ -409,7 +414,8 @@ def _create_architecture(
 
             id_to_node_list_map[layer.node_id] = layer_nodes
             nodes.extend(layer_nodes)
-            column_labels.append((str(layer.output_shape), current_x + node_size / 2, current_y))
+            label = format_shape_label(layer.output_shape, layer.extra_output_shapes)
+            column_labels.append((label, current_x + node_size / 2, current_y))
             current_y += 2 * node_size
 
         layer_y.append(current_y - node_spacing - 2 * node_size)

{visualtorch-1.1.0 → visualtorch-1.2.0}/visualtorch/lenet_style.py

@@ -17,7 +17,16 @@ from .backend import Architecture, extract_architecture
 from .connectors import compute_skip_levels, draw_connector
 from .utils.layer_utils import Input
 from .utils.traced_layer import TracedLayer
-from .utils.utils import ColorWheel, ImageDraw, InputShape, StackedBox, get_rgba_tuple, self_multiply
+from .utils.utils import (
+    ColorWheel,
+    ImageDraw,
+    InputShape,
+    StackedBox,
+    format_shape_label,
+    get_rgba_tuple,
+    resolve_palette,
+    self_multiply,
+)
 
 _LABEL_ROW_HEIGHT = 100
 
@@ -33,7 +42,8 @@ def lenet_view(
     scale_xy: float = 1,
     type_ignore: list | None = None,
     color_map: dict | None = None,
-    one_dim_orientation: str = "z",
+    palette: str = "okabe_ito",
+    low_dim_orientation: str = "z",
     background_fill: str | tuple[int, ...] = "white",
     padding: int = 10,
     spacing: int = 10,
@@ -69,7 +79,13 @@ def lenet_view(
         type_ignore (list, optional): List of layer types in the torch model to ignore during drawing.
         color_map (dict, optional): Dictionary defining fill and outline colors for each layer by class type.
             Will fallback to default values for unspecified classes.
-        one_dim_orientation (str, optional): Axis on which one-dim layers should be drawn. E.g., 'x', 'y', or 'z'.
+        palette (str, optional): Named color palette used as the fallback for any layer type not
+            given an explicit override via `color_map`. One of `"okabe_ito"` (default,
+            colorblind-safe), `"tol_bright"`, `"tol_muted"`, `"tab10"`, `"grayscale"`, `"nord"`,
+            `"dracula"`, `"gruvbox"`, `"solarized"`, `"material"`, `"catppuccin"`.
+        low_dim_orientation (str, optional): Axis on which a layer without real spatial/channel
+            structure (a 1D shape, or a 2D shape like an RNN/attention layer's
+            `(seq_len, hidden_size)`) should be drawn. One of `'x'`, `'y'`, or `'z'`.
         background_fill (str or tuple, optional): Background color for the image. A string or a tuple (R, G, B, A).
         padding (int, optional): Distance in pixels before the first and after the last layer.
         spacing (int, optional): Spacing in pixels between two layers.
@@ -122,7 +138,7 @@ def lenet_view(
 
     layer_types: list[type] = []
     make_box = _box_factory(
-        one_dim_orientation,
+        low_dim_orientation,
         scale_xy,
         min_xy,
         max_xy,
@@ -134,7 +150,7 @@ def lenet_view(
         opacity,
         offset_z,
         layer_types,
-        ColorWheel(),
+        ColorWheel(colors=resolve_palette(palette)),
     )
     column_layout = layout_columns(
         filtered_columns,
@@ -218,7 +234,7 @@ def _right_extent_for(box: VolumetricBox) -> float:
 
 
 def _box_factory(
-    one_dim_orientation: str,
+    low_dim_orientation: str,
     scale_xy: float,
     min_xy: int,
     max_xy: int,
@@ -237,19 +253,20 @@ def _box_factory(
     def make_box(layer: TracedLayer) -> StackedBox:
         shape = layer.output_shape[1:]  # drop batch size
 
-        if len(shape) == 1:
-            if one_dim_orientation in ("x", "y", "z"):
-                shape = (1,) * "cxyz".index(one_dim_orientation) + shape
+        if len(shape) in (1, 2):
+            # Neither a 1D nor a 2D shape has real spatial/channel structure - there's nothing
+            # to distinguish "channel" from "spatial" the way a genuine (C, H, W) feature map
+            # does. Take the last value (for 2D, e.g. an RNN/attention layer's
+            # (seq_len, hidden_size), this is the feature/channel-like one, matching PyTorch's
+            # (..., seq, feature) convention for sequence data; for 1D it's the only value) and
+            # let the user place it on whichever axis they choose, same as any 1D value - the
+            # positional-like dim, if any, is discarded either way.
+            value = shape[-1]
+            if low_dim_orientation in ("x", "y", "z"):
+                shape = (1,) * "cxyz".index(low_dim_orientation) + (value,)
             else:
-                error_msg = f"unsupported orientation: {one_dim_orientation}"
+                error_msg = f"unsupported orientation: {low_dim_orientation}"
                 raise ValueError(error_msg)
-        elif len(shape) == 2:
-            # A 2D non-batch shape (e.g. (seq_len, hidden_size) from an RNN/attention layer)
-            # isn't a CNN feature map missing a channel dim - there's no channel axis at all.
-            # StackedBox's slice count (de, below) is driven by shape[0], so a dummy 1 goes
-            # there instead of either real dim, keeping the two real dims on the box's actual
-            # width and height instead of one of them being drawn as that many stacked slices.
-            shape = (1, *shape)
 
         ori_shape = shape
         shape = shape + (1,) * (4 - len(shape))  # expand 4D.
@@ -266,6 +283,7 @@ def _box_factory(
         box.offset_z = offset_z
         box.label = layer.module.name() if isinstance(layer.module, Input) else layer_type.__name__
         box.output_shape = tuple(ori_shape)
+        box.extra_output_shapes = layer.extra_output_shapes
         box.de = z
 
         box.x1 = 0
@@ -405,6 +423,7 @@ def _draw_labels(
         for box in column:
             loc_x = box.x1 + (box.x2 - box.x1) // 4
             label = getattr(box, "label", type(box).__name__)
-            draw_text.text((loc_x, img.height - 50), f"{label} {box.output_shape}", font=font, fill=font_color)
+            shape_label = format_shape_label(box.output_shape, box.extra_output_shapes)
+            draw_text.text((loc_x, img.height - 50), f"{label} {shape_label}", font=font, fill=font_color)
 
     return Image.alpha_composite(img, text_img)

{visualtorch-1.1.0 → visualtorch-1.2.0}/visualtorch/render.py

@@ -32,6 +32,7 @@ class CommonOptions:
 
     to_file: str | None = None
     color_map: dict[Any, Any] | None = None
+    palette: str = "okabe_ito"
     background_fill: str | tuple[int, ...] = "white"
     padding: int = 10
     opacity: int = 255
@@ -66,7 +67,7 @@ class FlowStyleOptions:
     scale_z: float = 0.1
     scale_xy: float = 1
     type_ignore: list[type] | None = None
-    one_dim_orientation: str = "z"
+    low_dim_orientation: str = "z"
     draw_volume: bool = True
     spacing: int = 10
     draw_funnel: bool = True
@@ -86,7 +87,7 @@ class LenetStyleOptions:
     scale_z: float = 1
     scale_xy: float = 1
     type_ignore: list[type] | None = None
-    one_dim_orientation: str = "z"
+    low_dim_orientation: str = "z"
     spacing: int = 10
     draw_funnel: bool = True
     shade_step: int = 10
@@ -107,6 +108,7 @@ def _render_graph(
         input_shape,
         to_file=common.to_file,
         color_map=common.color_map,
+        palette=common.palette,
         node_size=options.node_size,
         background_fill=common.background_fill,
         padding=common.padding,
@@ -143,7 +145,8 @@ def _render_flow(
         scale_xy=options.scale_xy,
         type_ignore=options.type_ignore,
         color_map=common.color_map,
-        one_dim_orientation=options.one_dim_orientation,
+        palette=common.palette,
+        low_dim_orientation=options.low_dim_orientation,
         background_fill=common.background_fill,
         draw_volume=options.draw_volume,
         padding=common.padding,
@@ -177,7 +180,8 @@ def _render_lenet(
         scale_xy=options.scale_xy,
         type_ignore=options.type_ignore,
         color_map=common.color_map,
-        one_dim_orientation=options.one_dim_orientation,
+        palette=common.palette,
+        low_dim_orientation=options.low_dim_orientation,
         background_fill=common.background_fill,
         padding=common.padding,
         spacing=options.spacing,

{visualtorch-1.1.0 → visualtorch-1.2.0}/visualtorch/utils/recorder.py

@@ -106,28 +106,34 @@ def _wrap_and_stamp(obj: Any, ids: set[str]) -> Any:
     return obj
 
 
-def _first_tensor_shape(obj: Any) -> tuple[int, ...]:
-    """Recursively find the shape of the first tensor inside obj, or () if there isn't one."""
+def _all_tensor_shapes(obj: Any) -> list[tuple[int, ...]]:
+    """Recursively collect the shape of every tensor inside obj, in encounter order.
+
+    A module's return value isn't always a single tensor - `nn.LSTM` returns
+    `(output, (h_n, c_n))`, `nn.MultiheadAttention` returns `(attn_output, attn_weights)`, and a
+    custom module can return any tuple/list/dict of tensors. Collecting all of them (rather than
+    just the first) is what lets a caller show every output shape instead of silently dropping
+    every tensor after the first.
+    """
     if isinstance(obj, torch.Tensor):
-        return tuple(obj.shape)
+        return [tuple(obj.shape)]
     if isinstance(obj, Mapping):
+        shapes = []
         for value in obj.values():
-            shape = _first_tensor_shape(value)
-            if shape:
-                return shape
-        return ()
+            shapes.extend(_all_tensor_shapes(value))
+        return shapes
     if isinstance(obj, list | tuple):
+        shapes = []
         for value in obj:
-            shape = _first_tensor_shape(value)
-            if shape:
-                return shape
-        return ()
-    return ()
+            shapes.extend(_all_tensor_shapes(value))
+        return shapes
+    return []
 
 
 def _wrapped_module_call(
     id_to_module: dict[str, nn.Module],
     id_to_output_shape: dict[str, tuple[int, ...]],
+    id_to_extra_output_shapes: dict[str, tuple[tuple[int, ...], ...]],
     edges: list[tuple[str, str]],
     call_counts: dict[int, int],
 ) -> Any:
@@ -154,8 +160,11 @@ def _wrapped_module_call(
             call_counts[base_id] = call_index + 1
             node_id = f"{base_id}#{call_index}"
 
+            shapes = _all_tensor_shapes(out)
             id_to_module[node_id] = mod
-            id_to_output_shape[node_id] = _first_tensor_shape(out)
+            id_to_output_shape[node_id] = shapes[0] if shapes else ()
+            if len(shapes) > 1:
+                id_to_extra_output_shapes[node_id] = tuple(shapes[1:])
             edges.extend((producer_id, node_id) for producer_id in producer_ids)
             out = _wrap_and_stamp(out, {node_id})
 
@@ -171,11 +180,13 @@ class Recorder:
         self,
         id_to_module: dict[str, nn.Module],
         id_to_output_shape: dict[str, tuple[int, ...]],
+        id_to_extra_output_shapes: dict[str, tuple[tuple[int, ...], ...]],
         edges: list[tuple[str, str]],
         call_counts: dict[int, int],
     ) -> None:
         self._id_to_module = id_to_module
         self._id_to_output_shape = id_to_output_shape
+        self._id_to_extra_output_shapes = id_to_extra_output_shapes
         self._edges = edges
         self._call_counts = call_counts
 
@@ -184,6 +195,7 @@ class Recorder:
         nn.Module.__call__ = _wrapped_module_call(  # type: ignore[method-assign]
             self._id_to_module,
             self._id_to_output_shape,
+            self._id_to_extra_output_shapes,
             self._edges,
             self._call_counts,
         )
@@ -197,7 +209,13 @@ class Recorder:
 def trace_module_graph(
     model: nn.Module,
     input_shapes: tuple[tuple[int, ...], ...],
-) -> tuple[dict[str, nn.Module], dict[str, tuple[int, ...]], list[tuple[str, str]], list[str]]:
+) -> tuple[
+    dict[str, nn.Module],
+    dict[str, tuple[int, ...]],
+    dict[str, tuple[tuple[int, ...], ...]],
+    list[tuple[str, str]],
+    list[str],
+]:
     """Trace a forward pass to recover the leaf-module call graph.
 
     Args:
@@ -210,7 +228,11 @@ def trace_module_graph(
         tuple: A tuple containing:
             - id_to_module (dict): Mapping from node id to the leaf module. A leaf module
                 called more than once gets one entry per call (`f"{id(module)}#{call_index}"`).
-            - id_to_output_shape (dict): Mapping from node id to that module's output shape.
+            - id_to_output_shape (dict): Mapping from node id to that module's *first* output
+                tensor's shape (the one used for box sizing).
+            - id_to_extra_output_shapes (dict): Mapping from node id to the shapes of any
+                *additional* output tensors beyond the first (e.g. `nn.LSTM`'s `(h_n, c_n)`) -
+                only present for nodes that return more than one tensor.
             - edges (list): `(producer_node_id, consumer_node_id)` pairs, in call order.
             - input_ids (list): One synthetic node id per input tensor, in the same order as
                 `input_shapes`.
@@ -223,10 +245,11 @@ def trace_module_graph(
 
     id_to_module: dict[str, nn.Module] = {}
     id_to_output_shape: dict[str, tuple[int, ...]] = {}
+    id_to_extra_output_shapes: dict[str, tuple[tuple[int, ...], ...]] = {}
     edges: list[tuple[str, str]] = []
     call_counts: dict[int, int] = {}
 
-    with Recorder(id_to_module, id_to_output_shape, edges, call_counts):
+    with Recorder(id_to_module, id_to_output_shape, id_to_extra_output_shapes, edges, call_counts):
         if isinstance(model, nn.ModuleList):
             # nn.ModuleList has no forward() of its own - it's a plain container, not meant to
             # be called directly - so drive it the same way a user would: chain each child call.
@@ -244,4 +267,4 @@ def trace_module_graph(
             model(*dummy_inputs)
 
     input_ids = [f"{INPUT_NODE_ID}#{i}" for i in range(len(input_shapes))]
-    return id_to_module, id_to_output_shape, edges, input_ids
+    return id_to_module, id_to_output_shape, id_to_extra_output_shapes, edges, input_ids

{visualtorch-1.1.0 → visualtorch-1.2.0}/visualtorch/utils/traced_layer.py

@@ -20,3 +20,9 @@ class TracedLayer:
     module: nn.Module
     output_shape: tuple[int, ...]
     node_id: str
+    extra_output_shapes: tuple[tuple[int, ...], ...] = ()
+    """Shapes of any additional output tensors beyond `output_shape` (e.g. `nn.LSTM`'s hidden and
+    cell state, returned alongside its main sequence output) - empty for a module that returns
+    just one tensor. `output_shape` alone still drives box sizing; this is only ever read to
+    extend the `show_dimension` label so those extra tensors aren't silently unaccounted for.
+    """

{visualtorch-1.1.0 → visualtorch-1.2.0}/visualtorch/utils/utils.py

@@ -66,6 +66,7 @@ class StackedBox(Shape):
     offset_z: int
     label: str
     output_shape: tuple
+    extra_output_shapes: tuple[tuple[int, ...], ...] = ()
 
     def draw(self, draw: ImageDraw) -> None:
         """Draw box shape."""
@@ -114,6 +115,7 @@ class Box(Shape):
     de: int
     shade: int
     output_shape: tuple[int, ...]
+    extra_output_shapes: tuple[tuple[int, ...], ...] = ()
 
     def draw(self, draw: ImageDraw) -> None:
         """Draw box shape."""
@@ -225,26 +227,116 @@ class Ellipses(Shape):
         )
 
 
+PALETTES: dict[str, list[str]] = {
+    # Okabe-Ito: a colorblind-safe palette (Okabe & Ito, 2008) widely recommended for
+    # scientific visualization, e.g. in Nature's figure guidelines.
+    "okabe_ito": ["#E69F00", "#56B4E9", "#009E73", "#F0E442", "#0072B2", "#D55E00", "#CC79A7"],
+    # Paul Tol's "bright" qualitative scheme - colorblind-safe, higher-contrast alternative.
+    "tol_bright": ["#4477AA", "#EE6677", "#228833", "#CCBB44", "#66CCEE", "#AA3377", "#BBBBBB"],
+    # Paul Tol's "muted" qualitative scheme - colorblind-safe, softer aesthetic.
+    "tol_muted": [
+        "#CC6677",
+        "#332288",
+        "#DDCC77",
+        "#117733",
+        "#88CCEE",
+        "#882255",
+        "#44AA99",
+        "#999933",
+        "#AA4499",
+    ],
+    # matplotlib's default color cycle.
+    "tab10": [
+        "#1f77b4",
+        "#ff7f0e",
+        "#2ca02c",
+        "#d62728",
+        "#9467bd",
+        "#8c564b",
+        "#e377c2",
+        "#7f7f7f",
+        "#bcbd22",
+        "#17becf",
+    ],
+    # Evenly-spaced grays for print/monochrome-safe figures.
+    "grayscale": ["#404040", "#595959", "#737373", "#8c8c8c", "#a6a6a6", "#bfbfbf", "#d9d9d9"],
+    # Nord's Aurora + Frost accent colors.
+    "nord": [
+        "#bf616a",
+        "#d08770",
+        "#ebcb8b",
+        "#a3be8c",
+        "#b48ead",
+        "#8fbcbb",
+        "#88c0d0",
+        "#81a1c1",
+        "#5e81ac",
+    ],
+    # Dracula theme's accent colors.
+    "dracula": ["#FF5555", "#FFB86C", "#F1FA8C", "#50FA7B", "#8BE9FD", "#BD93F9", "#FF79C6"],
+    # Gruvbox's bright color variants.
+    "gruvbox": ["#fb4934", "#b8bb26", "#fabd2f", "#83a598", "#d3869b", "#8ec07c", "#fe8019"],
+    # Solarized's accent colors.
+    "solarized": [
+        "#b58900",
+        "#cb4b16",
+        "#dc322f",
+        "#d33682",
+        "#6c71c4",
+        "#268bd2",
+        "#2aa198",
+        "#859900",
+    ],
+    # Material Design's 500-weight color spread.
+    "material": [
+        "#f44336",
+        "#e91e63",
+        "#9c27b0",
+        "#3f51b5",
+        "#2196f3",
+        "#009688",
+        "#4caf50",
+        "#ffc107",
+        "#ff5722",
+    ],
+    # Catppuccin's Mocha flavor accent colors.
+    "catppuccin": [
+        "#f38ba8",
+        "#fab387",
+        "#f9e2af",
+        "#a6e3a1",
+        "#94e2d5",
+        "#89dceb",
+        "#89b4fa",
+        "#b4befe",
+        "#cba6f7",
+        "#f5c2e7",
+    ],
+}
+
+
+def resolve_palette(name: str) -> list[str]:
+    """Resolve a named palette to its list of hex colors.
+
+    Args:
+        name (str): One of the keys in `PALETTES`.
+
+    Returns:
+        list[str]: The palette's hex color strings.
+    """
+    if name not in PALETTES:
+        supported = ", ".join(sorted(PALETTES))
+        error_msg = f"Unsupported palette {name!r}. Supported palettes: {supported}."
+        raise ValueError(error_msg)
+    return PALETTES[name]
+
+
 class ColorWheel:
     """Default colors for the shapes."""
 
     def __init__(self, colors: list | None = None) -> None:
         self._cache: dict[type, Any] = {}
-        # Okabe-Ito: a colorblind-safe palette (Okabe & Ito, 2008) widely recommended for
-        # scientific visualization, e.g. in Nature's figure guidelines.
-        self.colors = (
-            colors
-            if colors is not None
-            else [
-                "#E69F00",  # orange
-                "#56B4E9",  # sky blue
-                "#009E73",  # bluish green
-                "#F0E442",  # yellow
-                "#0072B2",  # blue
-                "#D55E00",  # vermillion
-                "#CC79A7",  # reddish purple
-            ]
-        )
+        self.colors = colors if colors is not None else PALETTES["okabe_ito"]
 
     def get_color(self, class_type: type) -> tuple | None:
         """Return color from cache if exist, if not, get from the list and store it to the cache."""
@@ -344,6 +436,22 @@ def self_multiply(tensor_tuple: tuple | list) -> int | float:
     return s
 
 
+def format_shape_label(output_shape: tuple[int, ...], extra_output_shapes: tuple[tuple[int, ...], ...]) -> str:
+    """Format an output shape for display, appending any extra output shapes if present.
+
+    A module that returns more than one meaningful tensor (e.g. `nn.LSTM`'s `(output, (h_n,
+    c_n))`) would otherwise only ever show `output_shape` (the first tensor found, also the one
+    driving box size) with no indication the other tensors exist at all. `+` is used rather than
+    `visualtorch`'s existing `/` convention (already used to join sibling branches within one
+    column) so this doesn't read as an alternative/branch - these are all real outputs of this
+    same node, not one of several options.
+    """
+    label = str(output_shape)
+    if extra_output_shapes:
+        label += " + " + " + ".join(str(shape) for shape in extra_output_shapes)
+    return label
+
+
 def vertical_image_concat(
     im1: Image,
     im2: Image,

{visualtorch-1.1.0 → visualtorch-1.2.0/visualtorch.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: visualtorch
-Version: 1.1.0
+Version: 1.2.0
 Summary: Architecture visualization of Torch models
 Home-page: https://github.com/willyfh/visualtorch
 Author: Willy Fitra Hendria
@@ -54,21 +54,22 @@ Dynamic: summary
 
 </div>
 
-**VisualTorch** aims to help visualize Torch-based neural network architectures. It currently supports generating flow-style, graph-style, and LeNet-style architectures for PyTorch Sequential and Custom models. This tool is adapted from [visualkeras](https://github.com/paulgavrikov/visualkeras), [pytorchviz](https://github.com/szagoruyko/pytorchviz), and [pytorch-summary](https://github.com/sksq96/pytorch-summary).
+**VisualTorch** aims to help visualize Torch-based neural network architectures. It currently supports generating flow-style, graph-style, and LeNet-style architectures for PyTorch Sequential and Custom models. Its original visual styles were inspired by [visualkeras](https://github.com/paulgavrikov/visualkeras), [pytorchviz](https://github.com/szagoruyko/pytorchviz), [pytorch-summary](https://github.com/sksq96/pytorch-summary), and [torchview](https://github.com/mert-kurttutan/torchview); since then, it has grown its own unified tracing backend and architecture-handling logic well beyond its origins.
 
 **Note:** `1.0+` is a major release with breaking API changes, but with significantly better features and algorithms - upgrading is recommended. For the old API, use `0.2.5` or older.
 
-**Limitation:** VisualTorch traces a real forward pass to build the diagram, which has two inherent
-limitations shared by any tracing-based approach (not bugs, and not fixable without full symbolic
-execution): (1) models with **data-dependent control flow** (e.g. a branch only taken if a tensor
-value crosses some threshold) only show whichever branch the traced dummy input happened to take;
-(2) a layer that returns **multiple meaningful output tensors** (e.g. a custom multi-task head)
-only has its first tensor's shape reflected in that node's size/label - its downstream connections
-are still correct either way. Contributions are welcome!
+**Limitation:** VisualTorch traces a real forward pass to build the diagram, which has an inherent
+limitation shared by any tracing-based approach (not a bug, and not fixable without full symbolic
+execution): models with **data-dependent control flow** (e.g. a branch only taken if a tensor
+value crosses some threshold) only show whichever branch the traced dummy input happened to take.
+Separately, a layer that returns **multiple meaningful output tensors** (e.g. a custom multi-task
+head, or `nn.LSTM`'s `(output, (h_n, c_n))`) still has its node's size based on only its first
+tensor; with `show_dimension=True`, every output tensor's shape is shown in the label, not just
+the first. Downstream connections are correct either way. Contributions are welcome!
 
 <div align="center">
 
-![VisualTorch Examples](docs/source/_static/images/banners/readme-examples.png)
+![VisualTorch Examples](https://raw.githubusercontent.com/willyfh/visualtorch/e6ad79751e0f7412b1074beb45f9baeccd1419e4/docs/source/_static/images/banners/readme-examples.png)
 
 </div>
 
@@ -100,16 +101,16 @@ Please feel free to send a pull request to contribute to this project by followi
 
 This poject is available as open source under the terms of the [MIT License](https://github.com/willyfh/visualtorch/blob/main/LICENSE).
 
-Originally, this project was based on the [visualkeras](https://github.com/paulgavrikov/visualkeras) (under the MIT license), with additional modifications inspired by [pytorchviz](https://github.com/szagoruyko/pytorchviz), and [pytorch-summary](https://github.com/sksq96/pytorch-summary), both of which are also licensed under the MIT license.
+Originally, this project was based on the [visualkeras](https://github.com/paulgavrikov/visualkeras) (under the MIT license), with additional modifications inspired by [pytorchviz](https://github.com/szagoruyko/pytorchviz), [pytorch-summary](https://github.com/sksq96/pytorch-summary), and [torchview](https://github.com/mert-kurttutan/torchview), all of which are also licensed under the MIT license.
 
 ## Citation
 
 Please cite this project in your publications if it helps your research.
 
-**Note:** the paper below describes the API as of its publication date (2024). VisualTorch has
-since had breaking API changes (see the [documentation](https://visualtorch.readthedocs.io/en/latest/)
-for the current API) - the DOI always resolves to what was actually reviewed and published, so
-it isn't updated to match.
+**Note:** the paper below describes VisualTorch as of its publication date (2024). The project has
+since been substantially refactored, including breaking API changes (see the
+[documentation](https://visualtorch.readthedocs.io/en/latest/) for the current API) - the DOI
+always resolves to what was actually reviewed and published.
 
 ```bibtex
 @article{Hendria2024,