visualtorch 0.2.3__tar.gz → 0.2.5__tar.gz

{visualtorch-0.2.3/visualtorch.egg-info → visualtorch-0.2.5}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: visualtorch
-Version: 0.2.3
+Version: 0.2.5
 Summary: Architecture visualization of Torch models
 Home-page: https://github.com/willyfh/visualtorch
 Author: Willy Fitra Hendria
@@ -14,8 +14,37 @@ Classifier: License :: OSI Approved :: MIT License
 Classifier: Operating System :: OS Independent
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
-Provides-Extra: dev
 License-File: LICENSE
+Requires-Dist: pillow>=10.0.0
+Requires-Dist: numpy>=1.18.1
+Requires-Dist: aggdraw>=1.3.11
+Requires-Dist: torch>=2.0.0
+Provides-Extra: dev
+Requires-Dist: myst-parser; extra == "dev"
+Requires-Dist: nbsphinx; extra == "dev"
+Requires-Dist: pandoc; extra == "dev"
+Requires-Dist: sphinx<7.0; extra == "dev"
+Requires-Dist: sphinx_autodoc_typehints; extra == "dev"
+Requires-Dist: sphinx_book_theme; extra == "dev"
+Requires-Dist: sphinx-copybutton; extra == "dev"
+Requires-Dist: sphinx_design; extra == "dev"
+Requires-Dist: sphinx_gallery; extra == "dev"
+Requires-Dist: matplotlib; extra == "dev"
+Requires-Dist: pre-commit; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: keywords
+Dynamic: license
+Dynamic: license-file
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
 
 <div align="center">
  <h1>🔥 VisualTorch 🔥</h1>

{visualtorch-0.2.3 → visualtorch-0.2.5}/setup.py

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

visualtorch-0.2.5/tests/test_graph.py

@@ -0,0 +1,98 @@
+"""Tests for graph view."""
+
+# Copyright (C) 2024 Willy Fitra Hendria
+# SPDX-License-Identifier: MIT
+
+from pathlib import Path
+
+import pytest
+import torch
+from PIL import Image
+from torch import nn
+from visualtorch import graph_view
+
+
+@pytest.fixture()
+def dense_model() -> nn.Module:
+    """A simple dense model creation."""
+
+    class SimpleDense(nn.Module):
+        """Simple Dense Model."""
+
+        def __init__(self) -> None:
+            super().__init__()
+            self.h0 = nn.Linear(4, 8)
+            self.h1 = nn.Linear(8, 8)
+            self.h2 = nn.Linear(8, 4)
+            self.out = nn.Linear(4, 2)
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            """Define the forward pass."""
+            x = self.h0(x)
+            x = self.h1(x)
+            x = self.h2(x)
+            return self.out(x)
+
+    return SimpleDense()
+
+
+@pytest.fixture()
+def conv_model() -> nn.Module:
+    """A simple conv model, exercising the Conv2d/ConvolutionBackward0 path."""
+    return nn.Sequential(
+        nn.Conv2d(3, 8, 3, 1, 1),
+        nn.Conv2d(8, 16, 3, 1, 1),
+    )
+
+
+@pytest.fixture()
+def wide_dense_model() -> nn.Module:
+    """A dense model with a hidden layer wider than the default ellipsize_after threshold."""
+
+    class WideDense(nn.Module):
+        """A dense model with more than 10 hidden units, to trigger ellipsis drawing."""
+
+        def __init__(self) -> None:
+            super().__init__()
+            self.h0 = nn.Linear(4, 20)
+            self.out = nn.Linear(20, 2)
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            """Forward pass."""
+            return self.out(self.h0(x))
+
+    return WideDense()
+
+
+def test_dense_model_graph_view_runs(dense_model: nn.Module) -> None:
+    """Test graph view using dense model."""
+    _ = graph_view(dense_model, input_shape=(1, 4))
+
+
+def test_conv_model_graph_view_runs(conv_model: nn.Module) -> None:
+    """graph_view should support Conv2d layers, not just Linear."""
+    img = graph_view(conv_model, input_shape=(1, 3, 16, 16))
+    assert img is not None
+
+
+def test_graph_view_ellipsizes_wide_layers(wide_dense_model: nn.Module) -> None:
+    """A hidden layer wider than ellipsize_after should draw an ellipsis, not crash."""
+    img = graph_view(wide_dense_model, input_shape=(1, 4), ellipsize_after=10)
+    assert img is not None
+
+
+def test_graph_view_show_neurons_false(wide_dense_model: nn.Module) -> None:
+    """show_neurons=False should render one node per layer instead of per neuron."""
+    img = graph_view(wide_dense_model, input_shape=(1, 4), show_neurons=False)
+    assert img is not None
+
+
+def test_graph_view_writes_to_file(dense_model: nn.Module, tmp_path: Path) -> None:
+    """to_file should save a readable image to disk."""
+    out_file = tmp_path / "graph.png"
+    graph_view(dense_model, input_shape=(1, 4), to_file=str(out_file))
+
+    assert out_file.exists()
+    with Image.open(out_file) as saved_img:
+        assert saved_img.size[0] > 0
+        assert saved_img.size[1] > 0

visualtorch-0.2.5/tests/test_layered.py

@@ -0,0 +1,165 @@
+"""Tests for layered view."""
+
+# Copyright (C) 2024 Willy Fitra Hendria
+# SPDX-License-Identifier: MIT
+
+from pathlib import Path
+
+import pytest
+import torch
+import torch.nn.functional as func
+from PIL import Image
+from torch import nn
+from visualtorch import layered_view
+
+
+@pytest.fixture()
+def sequential_model() -> nn.Sequential:
+    """Define Sequential torch model for testing."""
+    return nn.Sequential(
+        nn.Conv2d(3, 64, 3, 1, 1),
+        nn.ReLU(),
+        nn.Conv2d(64, 128, 3, 1, 1),
+        nn.ReLU(),
+        nn.MaxPool2d(2, 2),
+    )
+
+
+@pytest.fixture()
+def module_list_model() -> nn.ModuleList:
+    """Define ModuleList-based torch model for testing."""
+    return nn.ModuleList(
+        [
+            nn.Conv2d(3, 64, 3, 1, 1),
+            nn.ReLU(),
+            nn.Conv2d(64, 128, 3, 1, 1),
+            nn.ReLU(),
+            nn.MaxPool2d(2, 2),
+        ],
+    )
+
+
+@pytest.fixture()
+def custom_model() -> nn.Module:
+    """Define the custom model."""
+
+    class CustomModel(nn.Module):
+        """A simple custom cnn model."""
+
+        def __init__(self) -> None:
+            super().__init__()
+            self.conv1 = nn.Conv2d(3, 64, 3, 1, 1)
+            self.conv2 = nn.Conv2d(64, 128, 3, 1, 1)
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            """Funcorward pass."""
+            x = func.relu(self.conv1(x))
+            x = func.relu(self.conv2(x))
+            return func.max_pool2d(x, 2, 2)
+
+    # Create an instance of the custom model
+    return CustomModel()
+
+
+@pytest.fixture()
+def lstm_model() -> nn.Module:
+    """Define a simple LSTM model for testing."""
+
+    class LSTMModel(nn.Module):
+        """A simple LSTM model."""
+
+        def __init__(self, input_size: int, hidden_size: int, num_layers: int) -> None:
+            super().__init__()
+            self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            """Forward pass."""
+            out, _ = self.lstm(x)
+            return out
+
+    # Create an instance of the LSTM model
+    return LSTMModel(input_size=10, hidden_size=20, num_layers=2)
+
+
+@pytest.fixture()
+def classifier_model() -> nn.Module:
+    """Define a model ending in a 1D (per-sample) output, e.g. classification logits."""
+
+    class ClassifierModel(nn.Module):
+        """A cnn model that ends with a 1D output."""
+
+        def __init__(self) -> None:
+            super().__init__()
+            self.conv = nn.Conv2d(3, 8, 3, 1, 1)
+            self.pool = nn.AdaptiveAvgPool2d(1)
+            self.fc = nn.Linear(8, 5)
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            """Forward pass."""
+            x = self.conv(x)
+            x = self.pool(x)
+            x = torch.flatten(x, 1)
+            return self.fc(x)
+
+    return ClassifierModel()
+
+
+def test_sequential_model_layered_view_runs(sequential_model: nn.Sequential) -> None:
+    """Test layered view on sequential model."""
+    _ = layered_view(sequential_model, input_shape=(1, 3, 224, 224))
+
+
+def test_module_list_model_layered_view_runs(module_list_model: nn.ModuleList) -> None:
+    """Test layered view on module list model."""
+    _ = layered_view(module_list_model, input_shape=(1, 3, 224, 224))
+
+
+def test_custom_model_layered_view_runs(custom_model: nn.Module) -> None:
+    """Test layered view on custom model."""
+    _ = layered_view(custom_model, input_shape=(1, 3, 224, 224))
+
+
+def test_lstm_model_layered_view_runs(lstm_model: nn.Module) -> None:
+    """Test layered view on lstm model."""
+    _ = layered_view(lstm_model, input_shape=(1, 10, 10))
+
+
+@pytest.mark.parametrize("orientation", ["x", "y", "z"])
+def test_layered_view_one_dim_orientation(classifier_model: nn.Module, orientation: str) -> None:
+    """Test layered view on a model with a 1D output, for every supported orientation."""
+    img = layered_view(classifier_model, input_shape=(1, 3, 16, 16), one_dim_orientation=orientation)
+    assert img is not None
+
+
+def test_layered_view_invalid_one_dim_orientation_raises(classifier_model: nn.Module) -> None:
+    """An unsupported one_dim_orientation should raise a clear ValueError."""
+    with pytest.raises(ValueError, match="unsupported orientation"):
+        layered_view(classifier_model, input_shape=(1, 3, 16, 16), one_dim_orientation="bad")
+
+
+def test_layered_view_with_type_and_index_ignore(sequential_model: nn.Sequential) -> None:
+    """Layers matched by type_ignore or index_ignore should be skipped without error."""
+    img = layered_view(
+        sequential_model,
+        input_shape=(1, 3, 224, 224),
+        type_ignore=[nn.ReLU],
+        index_ignore=[0],
+    )
+    assert img is not None
+
+
+def test_layered_view_with_legend(sequential_model: nn.Sequential) -> None:
+    """legend=True should append a legend without error."""
+    img = layered_view(sequential_model, input_shape=(1, 3, 224, 224), legend=True)
+    assert img is not None
+
+
+def test_layered_view_writes_to_file(sequential_model: nn.Sequential, tmp_path: Path) -> None:
+    """to_file should save a readable image to disk."""
+    out_file = tmp_path / "layered.png"
+    layered_view(sequential_model, input_shape=(1, 3, 224, 224), to_file=str(out_file))
+
+    assert out_file.exists()
+    with Image.open(out_file) as saved_img:
+        assert saved_img.size[0] > 0
+        assert saved_img.size[1] > 0

visualtorch-0.2.5/tests/test_lenet_style.py

@@ -0,0 +1,159 @@
+"""Tests for lenet view."""
+
+# Copyright (C) 2024 Willy Fitra Hendria
+# SPDX-License-Identifier: MIT
+
+from pathlib import Path
+
+import pytest
+import torch
+import torch.nn.functional as func
+from PIL import Image
+from torch import nn
+from visualtorch import lenet_view
+
+
+@pytest.fixture()
+def sequential_model() -> nn.Sequential:
+    """Define Sequential torch model for testing."""
+    return nn.Sequential(
+        nn.Conv2d(3, 64, 3, 1, 1),
+        nn.ReLU(),
+        nn.Conv2d(64, 128, 3, 1, 1),
+        nn.ReLU(),
+        nn.MaxPool2d(2, 2),
+    )
+
+
+@pytest.fixture()
+def module_list_model() -> nn.ModuleList:
+    """Define ModuleList-based torch model for testing."""
+    return nn.ModuleList(
+        [
+            nn.Conv2d(3, 64, 3, 1, 1),
+            nn.ReLU(),
+            nn.Conv2d(64, 128, 3, 1, 1),
+            nn.ReLU(),
+            nn.MaxPool2d(2, 2),
+        ],
+    )
+
+
+@pytest.fixture()
+def custom_model() -> nn.Module:
+    """Define the custom model."""
+
+    class CustomModel(nn.Module):
+        """A simple custom cnn model."""
+
+        def __init__(self) -> None:
+            super().__init__()
+            self.conv1 = nn.Conv2d(3, 64, 3, 1, 1)
+            self.conv2 = nn.Conv2d(64, 128, 3, 1, 1)
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            """Funcorward pass."""
+            x = func.relu(self.conv1(x))
+            x = func.relu(self.conv2(x))
+            return func.max_pool2d(x, 2, 2)
+
+    # Create an instance of the custom model
+    return CustomModel()
+
+
+@pytest.fixture()
+def lstm_model() -> nn.Module:
+    """Define a simple LSTM model for testing."""
+
+    class LSTMModel(nn.Module):
+        """A simple LSTM model."""
+
+        def __init__(self, input_size: int, hidden_size: int, num_layers: int) -> None:
+            super().__init__()
+            self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            """Forward pass."""
+            out, _ = self.lstm(x)
+            return out
+
+    # Create an instance of the LSTM model
+    return LSTMModel(input_size=10, hidden_size=20, num_layers=2)
+
+
+@pytest.fixture()
+def classifier_model() -> nn.Module:
+    """Define a model ending in a 1D (per-sample) output, e.g. classification logits."""
+
+    class ClassifierModel(nn.Module):
+        """A cnn model that ends with a 1D output."""
+
+        def __init__(self) -> None:
+            super().__init__()
+            self.conv = nn.Conv2d(3, 8, 3, 1, 1)
+            self.pool = nn.AdaptiveAvgPool2d(1)
+            self.fc = nn.Linear(8, 5)
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            """Forward pass."""
+            x = self.conv(x)
+            x = self.pool(x)
+            x = torch.flatten(x, 1)
+            return self.fc(x)
+
+    return ClassifierModel()
+
+
+def test_sequential_model_lenet_view_runs(sequential_model: nn.Sequential) -> None:
+    """Test lenet view on sequential model."""
+    _ = lenet_view(sequential_model, input_shape=(1, 3, 224, 224))
+
+
+def test_module_list_model_lenet_view_runs(module_list_model: nn.ModuleList) -> None:
+    """Test lenet view on module list model."""
+    _ = lenet_view(module_list_model, input_shape=(1, 3, 224, 224))
+
+
+def test_custom_model_lenet_view_runs(custom_model: nn.Module) -> None:
+    """Test lenet view on custom model."""
+    _ = lenet_view(custom_model, input_shape=(1, 3, 224, 224))
+
+
+def test_lstm_model_layered_view_runs(lstm_model: nn.Module) -> None:
+    """Test layered view on lstm model."""
+    _ = lenet_view(lstm_model, input_shape=(1, 10, 10))
+
+
+@pytest.mark.parametrize("orientation", ["x", "y", "z"])
+def test_lenet_view_one_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)
+    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."""
+    with pytest.raises(ValueError, match="unsupported orientation"):
+        lenet_view(classifier_model, input_shape=(1, 3, 16, 16), one_dim_orientation="bad")
+
+
+def test_lenet_view_with_type_and_index_ignore(sequential_model: nn.Sequential) -> None:
+    """Layers matched by type_ignore or index_ignore should be skipped without error."""
+    img = lenet_view(
+        sequential_model,
+        input_shape=(1, 3, 224, 224),
+        type_ignore=[nn.ReLU],
+        index_ignore=[0],
+    )
+    assert img is not None
+
+
+def test_lenet_view_writes_to_file(sequential_model: nn.Sequential, tmp_path: Path) -> None:
+    """to_file should save a readable image to disk."""
+    out_file = tmp_path / "lenet.png"
+    lenet_view(sequential_model, input_shape=(1, 3, 224, 224), to_file=str(out_file))
+
+    assert out_file.exists()
+    with Image.open(out_file) as saved_img:
+        assert saved_img.size[0] > 0
+        assert saved_img.size[1] > 0

visualtorch-0.2.5/tests/test_regression_issues.py

@@ -0,0 +1,73 @@
+"""Regression tests for shape-handling crashes reported in open GitHub issues.
+
+See https://github.com/willyfh/visualtorch/issues/63, /68, /69.
+"""
+
+# Copyright (C) 2024 Willy Fitra Hendria
+# SPDX-License-Identifier: MIT
+
+import pytest
+import torch
+from torch import nn
+from visualtorch import layered_view, lenet_view
+from visualtorch.utils.utils import self_multiply
+
+
+@pytest.fixture()
+def multi_output_container_model() -> nn.Module:
+    """A model whose inner block is a container (not Sequential/ModuleList) that returns multiple tensors.
+
+    This mirrors timm's ``FeatureListNet`` used in issue #69: a container module that isn't
+    ``nn.Sequential``/``nn.ModuleList`` and returns multiple differently-shaped tensors.
+    """
+
+    class FeaturePyramidBlock(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.stage1 = nn.Conv2d(3, 8, kernel_size=3, stride=2, padding=1)
+            self.stage2 = nn.Conv2d(8, 16, kernel_size=3, stride=2, padding=1)
+            self.stage3 = nn.Conv2d(16, 32, kernel_size=3, stride=2, padding=1)
+
+        def forward(self, x: torch.Tensor) -> list:
+            f1 = self.stage1(x)
+            f2 = self.stage2(f1)
+            f3 = self.stage3(f2)
+            return [f1, f2, f3]
+
+    class MultiScaleNet(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.features = FeaturePyramidBlock()
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            return self.features(x)[-1]
+
+    return MultiScaleNet()
+
+
+def test_layered_view_multi_output_container(multi_output_container_model: nn.Module) -> None:
+    """layered_view should not crash on container modules that output multiple tensors."""
+    img = layered_view(multi_output_container_model, input_shape=(1, 3, 32, 32))
+    assert img is not None
+
+
+def test_lenet_view_multi_output_container(multi_output_container_model: nn.Module) -> None:
+    """lenet_view should not crash on container modules that output multiple tensors."""
+    img = lenet_view(multi_output_container_model, input_shape=(1, 3, 32, 32))
+    assert img is not None
+
+
+def test_self_multiply_handles_nested_shape() -> None:
+    """self_multiply should always reduce to a scalar, even if an element is itself a shape."""
+    nested_shape = (1, torch.Size([4, 8]))
+    result = self_multiply(nested_shape)
+    assert isinstance(result, int)
+
+
+def test_lenet_view_rejects_multi_input_shape_with_clear_error() -> None:
+    """Multi-tensor-input models aren't supported yet; the failure should be a clear ValueError."""
+    model = nn.Linear(10, 5)
+    multi_input_shape = ((1, 10), (1, 10))
+
+    with pytest.raises(ValueError, match="single"):
+        lenet_view(model, input_shape=multi_input_shape)

{visualtorch-0.2.3 → visualtorch-0.2.5}/visualtorch/layered.py

@@ -22,6 +22,7 @@ from .utils.utils import (
     get_rgba_tuple,
     linear_layout,
     self_multiply,
+    validate_input_shape,
     vertical_image_concat,
 )
 
@@ -85,6 +86,8 @@ def layered_view(
     """
     # Iterate over the model to compute bounds and generate boxes
 
+    validate_input_shape(input_shape)
+
     x_off = -1
 
     img_height = 0
@@ -291,7 +294,7 @@ def _create_architecture(
         layer = layers[key]["module"]
         shape = layers[key]["output_shape"]
         # Do no render the SpacingDummyLayer, just increase the pointer
-        if type(layer) == SpacingDummyLayer:
+        if type(layer) is SpacingDummyLayer:
             current_z += layer.spacing
             continue
 

{visualtorch-0.2.3 → visualtorch-0.2.5}/visualtorch/lenet_style.py

@@ -19,6 +19,7 @@ from .utils.utils import (
     StackedBox,
     get_rgba_tuple,
     self_multiply,
+    validate_input_shape,
 )
 
 
@@ -82,6 +83,8 @@ def lenet_view(
     """
     # Iterate over the model to compute bounds and generate boxes
 
+    validate_input_shape(input_shape)
+
     x_off = -1
 
     img_height = 0
@@ -247,7 +250,7 @@ def _create_architecture(
         layer = layers[key]["module"]
         shape = layers[key]["output_shape"]
         # Do no render the SpacingDummyLayer, just increase the pointer
-        if type(layer) == SpacingDummyLayer:
+        if type(layer) is SpacingDummyLayer:
             current_x += layer.spacing
             continue
 

{visualtorch-0.2.3 → visualtorch-0.2.5}/visualtorch/utils/layer_utils.py

@@ -185,12 +185,20 @@ def register_hook(
         m_key = "%s-%i" % (class_name, module_idx + 1)
         layers[m_key] = OrderedDict()
         layers[m_key]["module"] = module
-        if isinstance(out, list | tuple):
-            layers[m_key]["output_shape"] = tuple((-1,) + o.size()[1:] for o in out)
+        if isinstance(out, tuple | list):
+            if len(out) > 0 and hasattr(out[0], "size"):
+                layers[m_key]["output_shape"] = out[0].size()
+            else:
+                layers[m_key]["output_shape"] = tuple(o.size() for o in out if hasattr(o, "size"))
         else:
             layers[m_key]["output_shape"] = out.size()
 
-    if not isinstance(module, nn.Sequential) and not isinstance(module, nn.ModuleList) and module is not model:
+    # Only hook leaf modules (no children). Container modules - whether nn.Sequential,
+    # nn.ModuleList, or a custom container such as timm's FeatureListNet - would otherwise be
+    # captured as if they were a single layer, with their multi-tensor output mistaken for one
+    # layer's output shape.
+    is_leaf = len(list(module.children())) == 0
+    if is_leaf and module is not model:
         hooks.append(module.register_forward_hook(hook))
 
 

{visualtorch-0.2.3 → visualtorch-0.2.5}/visualtorch/utils/utils.py

@@ -261,23 +261,44 @@ def get_keys_by_value(d: dict, v: int) -> Generator:
             yield key
 
 
-def self_multiply(tensor_tuple: tuple) -> int | float:
+def validate_input_shape(input_shape: tuple) -> None:
+    """Validate that input_shape describes a single input tensor.
+
+    Args:
+        input_shape (tuple): The shape to validate.
+
+    Raises:
+        ValueError: If input_shape is not a single flat tuple of ints, e.g. when a tuple of
+            per-tensor shapes was passed for a model that takes multiple separate input tensors.
+    """
+    if not isinstance(input_shape, tuple) or not all(isinstance(dim, int) for dim in input_shape):
+        error_msg = (
+            "input_shape must be a single tuple of ints, e.g. (1, 3, 224, 224). "
+            f"Got {input_shape!r} instead. Visualizing models that take multiple separate "
+            "input tensors is not supported yet."
+        )
+        raise ValueError(error_msg)
+
+
+def self_multiply(tensor_tuple: tuple | list) -> int | float:
     """Multiplies all elements in the tuple together.
 
+    Elements that are themselves a tuple/list (e.g. a nested torch.Size, which can end up here
+    when a layer's captured output shape wasn't a plain flat shape) are flattened by multiplying
+    their own elements together first, so the result is always a scalar.
+
     Args:
-        tensor_tuple (tuple): A tuple containing tensors.
+        tensor_tuple (tuple or list): A tuple containing tensors.
 
     Returns:
         int or float: The result of multiplying all elements together.
     """
-    tensor_list = list(tensor_tuple)
-    if None in tensor_list:
-        tensor_list.remove(None)
+    tensor_list = [v for v in tensor_tuple if v is not None]
     if len(tensor_list) == 0:
         return 0
-    s = tensor_list[0]
-    for i in range(1, len(tensor_list)):
-        s *= tensor_list[i]
+    s = 1
+    for v in tensor_list:
+        s *= self_multiply(v) if isinstance(v, tuple | list) else v
     return s
 
 

{visualtorch-0.2.3 → visualtorch-0.2.5/visualtorch.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: visualtorch
-Version: 0.2.3
+Version: 0.2.5
 Summary: Architecture visualization of Torch models
 Home-page: https://github.com/willyfh/visualtorch
 Author: Willy Fitra Hendria
@@ -14,8 +14,37 @@ Classifier: License :: OSI Approved :: MIT License
 Classifier: Operating System :: OS Independent
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
-Provides-Extra: dev
 License-File: LICENSE
+Requires-Dist: pillow>=10.0.0
+Requires-Dist: numpy>=1.18.1
+Requires-Dist: aggdraw>=1.3.11
+Requires-Dist: torch>=2.0.0
+Provides-Extra: dev
+Requires-Dist: myst-parser; extra == "dev"
+Requires-Dist: nbsphinx; extra == "dev"
+Requires-Dist: pandoc; extra == "dev"
+Requires-Dist: sphinx<7.0; extra == "dev"
+Requires-Dist: sphinx_autodoc_typehints; extra == "dev"
+Requires-Dist: sphinx_book_theme; extra == "dev"
+Requires-Dist: sphinx-copybutton; extra == "dev"
+Requires-Dist: sphinx_design; extra == "dev"
+Requires-Dist: sphinx_gallery; extra == "dev"
+Requires-Dist: matplotlib; extra == "dev"
+Requires-Dist: pre-commit; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: keywords
+Dynamic: license
+Dynamic: license-file
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
 
 <div align="center">
  <h1>🔥 VisualTorch 🔥</h1>

{visualtorch-0.2.3 → visualtorch-0.2.5}/visualtorch.egg-info/SOURCES.txt

@@ -2,6 +2,10 @@ LICENSE
 README.md
 pyproject.toml
 setup.py
+tests/test_graph.py
+tests/test_layered.py
+tests/test_lenet_style.py
+tests/test_regression_issues.py
 visualtorch/__init__.py
 visualtorch/graph.py
 visualtorch/layered.py