lucid-dl 2.11.3__tar.gz → 2.11.5__tar.gz

{lucid_dl-2.11.3/lucid_dl.egg-info → lucid_dl-2.11.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lucid-dl
-Version: 2.11.3
+Version: 2.11.5
 Summary: Lumerico's Comprehensive Interface for Deep Learning
 Home-page: https://github.com/ChanLumerico/lucid
 Author: ChanLumerico
@@ -48,19 +48,9 @@ Whether you're a student, educator, or an advanced researcher seeking to demysti
 
 ### 🔥 What's New
 
-- Added additional `nn.Module` hooks for richer introspection during training:
-
-  ```python
-  def register_forward_pre_hook(self, hook: Callable, *, with_kwargs: bool = False)
-
-  def register_forward_hook(self, hook: Callable, *, with_kwargs: bool = False)
+- Added new visual tool: `lucid.visual.build_tensor_mermaid_chart` which builds a Mermaid chart of given tensor's computatoinal graph
 
-  def register_backward_hook(self, hook: Callable)
-
-  def register_full_backward_pre_hook(self, hook: Callable)
-
-  def register_full_backward_hook(self, hook: Callable)
-  ```
+- Added additional `nn.Module` hooks for richer introspection during training:
 
 ## 🔧 How to Install
 

{lucid_dl-2.11.3 → lucid_dl-2.11.5}/README.md

@@ -20,19 +20,9 @@ Whether you're a student, educator, or an advanced researcher seeking to demysti
 
 ### 🔥 What's New
 
-- Added additional `nn.Module` hooks for richer introspection during training:
-
-  ```python
-  def register_forward_pre_hook(self, hook: Callable, *, with_kwargs: bool = False)
-
-  def register_forward_hook(self, hook: Callable, *, with_kwargs: bool = False)
+- Added new visual tool: `lucid.visual.build_tensor_mermaid_chart` which builds a Mermaid chart of given tensor's computatoinal graph
 
-  def register_backward_hook(self, hook: Callable)
-
-  def register_full_backward_pre_hook(self, hook: Callable)
-
-  def register_full_backward_hook(self, hook: Callable)
-  ```
+- Added additional `nn.Module` hooks for richer introspection during training:
 
 ## 🔧 How to Install
 

{lucid_dl-2.11.3 → lucid_dl-2.11.5}/lucid/nn/module.py

@@ -13,7 +13,22 @@ from typing import (
 from collections import OrderedDict
 
 from lucid._tensor import Tensor
-from lucid.types import _ArrayOrScalar, _NumPyArray, _DeviceType
+from lucid.types import (
+    _ArrayOrScalar,
+    _BackwardHook,
+    _DeviceType,
+    _ForwardHook,
+    _ForwardHookKwargs,
+    _ForwardPreHook,
+    _ForwardPreHookKwargs,
+    _FullBackwardHook,
+    _FullBackwardPreHook,
+    _LoadStateDictPostHook,
+    _LoadStateDictPreHook,
+    _NumPyArray,
+    _StateDictHook,
+    _StateDictPreHook,
+)
 
 import lucid.nn as nn
 
@@ -30,26 +45,6 @@ __all__ = [
 ]
 
 
-_ForwardPreHook = Callable[["Module", tuple[Any, ...]], tuple[Any, ...] | None]
-_ForwardPreHookKwargs = Callable[
-    ["Module", tuple[Any, ...], dict[str, Any]],
-    tuple[tuple[Any, ...], dict[str, Any]] | None,
-]
-_ForwardHook = Callable[["Module", tuple[Any, ...], Any], Any | None]
-_ForwardHookKwargs = Callable[
-    ["Module", tuple[Any, ...], dict[str, Any], Any], Any | None
-]
-
-_BackwardHook = Callable[[Tensor, _NumPyArray], None]
-_FullBackwardPreHook = Callable[
-    ["Module", tuple[_NumPyArray | None, ...]], tuple[_NumPyArray | None, ...] | None
-]
-_FullBackwardHook = Callable[
-    ["Module", tuple[_NumPyArray | None, ...], tuple[_NumPyArray | None, ...]],
-    tuple[_NumPyArray | None, ...] | None,
-]
-
-
 class Module:
     _registry_map: dict[Type, OrderedDict[str, Any]] = {}
     _alt_name: str = ""
@@ -70,10 +65,17 @@ class Module:
             tuple[_ForwardPreHook | _ForwardPreHookKwargs, bool]
         ] = []
         self._forward_hooks: list[tuple[_ForwardHook | _ForwardHookKwargs, bool]] = []
+
         self._backward_hooks: list[_BackwardHook] = []
         self._full_backward_pre_hooks: list[_FullBackwardPreHook] = []
         self._full_backward_hooks: list[_FullBackwardHook] = []
 
+        self._state_dict_pre_hooks: list[_StateDictPreHook] = []
+        self._state_dict_hooks: list[_StateDictHook] = []
+
+        self._load_state_dict_pre_hooks: list[_LoadStateDictPreHook] = []
+        self._load_state_dict_post_hooks: list[_LoadStateDictPostHook] = []
+
         self._state_dict_pass_attr = set()
 
     def __setattr__(self, name: str, value: Any) -> None:
@@ -155,6 +157,26 @@ class Module:
         self._full_backward_hooks.append(hook)
         return lambda: self._full_backward_hooks.remove(hook)
 
+    def register_state_dict_pre_hook(self, hook: _StateDictPreHook) -> Callable:
+        self._state_dict_pre_hooks.append(hook)
+        return lambda: self._state_dict_pre_hooks.remove(hook)
+
+    def register_state_dict_hook(self, hook: _StateDictHook) -> Callable:
+        self._state_dict_hooks.append(hook)
+        return lambda: self._state_dict_hooks.remove(hook)
+
+    def register_load_state_dict_pre_hook(
+        self, hook: _LoadStateDictPreHook
+    ) -> Callable:
+        self._load_state_dict_pre_hooks.append(hook)
+        return lambda: self._load_state_dict_pre_hooks.remove(hook)
+
+    def register_load_state_dict_post_hook(
+        self, hook: _LoadStateDictPostHook
+    ) -> Callable:
+        self._load_state_dict_post_hooks.append(hook)
+        return lambda: self._load_state_dict_post_hooks.remove(hook)
+
     def reset_parameters(self) -> None:
         for param in self.parameters():
             param.zero()
@@ -231,6 +253,9 @@ class Module:
         prefix: str = "",
         keep_vars: bool = False,
     ) -> OrderedDict:
+        for hook in self._state_dict_pre_hooks:
+            hook(self, prefix, keep_vars)
+
         if destination is None:
             destination = OrderedDict()
 
@@ -249,9 +274,15 @@ class Module:
             if key in self._state_dict_pass_attr:
                 del destination[key]
 
+        for hook in self._state_dict_hooks:
+            hook(self, destination, prefix, keep_vars)
+
         return destination
 
     def load_state_dict(self, state_dict: OrderedDict, strict: bool = True) -> None:
+        for hook in self._load_state_dict_pre_hooks:
+            hook(self, state_dict, strict)
+
         own_state = self.state_dict(keep_vars=True)
 
         missing_keys = set(own_state.keys()) - set(state_dict.keys())
@@ -277,6 +308,9 @@ class Module:
             elif strict:
                 raise KeyError(f"Unexpected key '{key}' in state_dict.")
 
+        for hook in self._load_state_dict_post_hooks:
+            hook(self, missing_keys, unexpected_keys, strict)
+
     def __call__(self, *args: Any, **kwargs: Any) -> Tensor | tuple[Tensor, ...]:
         for hook, with_kwargs in self._forward_pre_hooks:
             if with_kwargs:

{lucid_dl-2.11.3 → lucid_dl-2.11.5}/lucid/types.py

@@ -6,8 +6,10 @@ from typing import (
     Sequence,
     Literal,
     TypeAlias,
+    TYPE_CHECKING,
     runtime_checkable,
 )
+from collections import OrderedDict
 import re
 
 import numpy as np
@@ -76,6 +78,62 @@ class _TensorLike(Protocol):
     ) -> None: ...
 
 
+@runtime_checkable
+class _ModuleHookable(Protocol):
+    def register_forward_pre_hook(
+        self, hook: Callable, *, with_kwargs: bool = False
+    ) -> Callable: ...
+
+    def register_forward_hook(
+        self, hook: Callable, *, with_kwargs: bool = False
+    ) -> Callable: ...
+
+    def register_backward_hook(self, hook: Callable) -> Callable: ...
+
+    def register_full_backward_pre_hook(self, hook: Callable) -> Callable: ...
+
+    def register_full_backward_hook(self, hook: Callable) -> Callable: ...
+
+    def register_state_dict_pre_hook(self, hook: Callable) -> Callable: ...
+
+    def register_state_dict_hook(self, hook: Callable) -> Callable: ...
+
+    def register_load_state_dict_pre_hook(self, hook: Callable) -> Callable: ...
+
+    def register_load_state_dict_post_hook(self, hook: Callable) -> Callable: ...
+
+
+_ForwardPreHook: TypeAlias = Callable[
+    [_ModuleHookable, tuple[Any, ...]], tuple[Any, ...] | None
+]
+_ForwardPreHookKwargs: TypeAlias = Callable[
+    [_ModuleHookable, tuple[Any, ...], dict[str, Any]],
+    tuple[tuple[Any, ...], dict[str, Any]] | None,
+]
+_ForwardHook: TypeAlias = Callable[[_ModuleHookable, tuple[Any, ...], Any], Any | None]
+_ForwardHookKwargs: TypeAlias = Callable[
+    [_ModuleHookable, tuple[Any, ...], dict[str, Any], Any], Any | None
+]
+
+_BackwardHook: TypeAlias = Callable[[_TensorLike, _NumPyArray], None]
+_FullBackwardPreHook: TypeAlias = Callable[
+    [_ModuleHookable, tuple[_NumPyArray | None, ...]],
+    tuple[_NumPyArray | None, ...] | None,
+]
+_FullBackwardHook: TypeAlias = Callable[
+    [_ModuleHookable, tuple[_NumPyArray | None, ...], tuple[_NumPyArray | None, ...]],
+    tuple[_NumPyArray | None, ...] | None,
+]
+
+_StateDictPreHook: TypeAlias = Callable[[_ModuleHookable, str, bool], None]
+_StateDictHook: TypeAlias = Callable[[_ModuleHookable, OrderedDict, str, bool], None]
+
+_LoadStateDictPreHook: TypeAlias = Callable[[_ModuleHookable, OrderedDict, bool], None]
+_LoadStateDictPostHook: TypeAlias = Callable[
+    [_ModuleHookable, set[str], set[str], bool], None
+]
+
+
 class Numeric:
     def __init__(
         self, base_dtype: type[int | float | complex], bits: int | None

{lucid_dl-2.11.3 → lucid_dl-2.11.5}/lucid/visual/__init__.py

@@ -1,2 +1 @@
-from .graph import *
 from .mermaid import *

{lucid_dl-2.11.3 → lucid_dl-2.11.5}/lucid/visual/mermaid.py

@@ -9,7 +9,7 @@ from lucid._tensor import Tensor
 from lucid.types import _ShapeLike
 
 
-__all__ = ["build_mermaid_chart"]
+__all__ = ["build_tensor_mermaid_chart", "build_module_mermaid_chart"]
 
 
 _NN_MODULES_PREFIX = "lucid.nn.modules."
@@ -255,7 +255,7 @@ def _collapse_repeated_children(
     return out
 
 
-def build_mermaid_chart(
+def build_module_mermaid_chart(
     module: nn.Module,
     input_shape: _ShapeLike | list[_ShapeLike] | None = None,
     inputs: Iterable[Tensor] | Tensor | None = None,
@@ -751,6 +751,192 @@ def build_mermaid_chart(
     return text
 
 
+def build_mermaid_chart(
+    module: nn.Module,
+    input_shape: _ShapeLike | list[_ShapeLike] | None = None,
+    inputs: Iterable[Tensor] | Tensor | None = None,
+    depth: int = 2,
+    direction: str = "LR",
+    include_io: bool = True,
+    show_params: bool = False,
+    return_lines: bool = False,
+    copy_to_clipboard: bool = False,
+    compact: bool = False,
+    use_class_defs: bool = False,
+    end_semicolons: bool = True,
+    edge_mode: Literal["dataflow", "execution"] = "execution",
+    collapse_repeats: bool = True,
+    repeat_min: int = 2,
+    color_by_subpackage: bool = True,
+    container_name_from_attr: bool = True,
+    edge_stroke_width: float = 2.0,
+    emphasize_model_title: bool = True,
+    model_title_font_px: int = 20,
+    show_shapes: bool = False,
+    hide_subpackages: Iterable[str] = (),
+    hide_module_names: Iterable[str] = (),
+    dash_multi_input_edges: bool = True,
+    subgraph_fill: str = "#000000",
+    subgraph_fill_opacity: float = 0.05,
+    subgraph_stroke: str = "#000000",
+    subgraph_stroke_opacity: float = 0.75,
+    force_text_color: str | None = None,
+    edge_curve: str = "natural",
+    node_spacing: int = 50,
+    rank_spacing: int = 50,
+    **forward_kwargs,
+) -> str | list[str]:
+    return build_module_mermaid_chart(
+        module,
+        input_shape=input_shape,
+        inputs=inputs,
+        depth=depth,
+        direction=direction,
+        include_io=include_io,
+        show_params=show_params,
+        return_lines=return_lines,
+        copy_to_clipboard=copy_to_clipboard,
+        compact=compact,
+        use_class_defs=use_class_defs,
+        end_semicolons=end_semicolons,
+        edge_mode=edge_mode,
+        collapse_repeats=collapse_repeats,
+        repeat_min=repeat_min,
+        color_by_subpackage=color_by_subpackage,
+        container_name_from_attr=container_name_from_attr,
+        edge_stroke_width=edge_stroke_width,
+        emphasize_model_title=emphasize_model_title,
+        model_title_font_px=model_title_font_px,
+        show_shapes=show_shapes,
+        hide_subpackages=hide_subpackages,
+        hide_module_names=hide_module_names,
+        dash_multi_input_edges=dash_multi_input_edges,
+        subgraph_fill=subgraph_fill,
+        subgraph_fill_opacity=subgraph_fill_opacity,
+        subgraph_stroke=subgraph_stroke,
+        subgraph_stroke_opacity=subgraph_stroke_opacity,
+        force_text_color=force_text_color,
+        edge_curve=edge_curve,
+        node_spacing=node_spacing,
+        rank_spacing=rank_spacing,
+        **forward_kwargs,
+    )
+
+
+def build_tensor_mermaid_chart(
+    tensor: Tensor,
+    horizontal: bool = False,
+    title: str | None = None,
+    start_id: int | None = None,
+    end_semicolons: bool = True,
+    copy_to_clipboard: bool = False,
+    use_class_defs: bool = True,
+    op_fill: str = "lightgreen",
+    param_fill: str = "plum",
+    result_fill: str = "lightcoral",
+    leaf_fill: str = "lightgray",
+    grad_fill: str = "lightblue",
+    start_fill: str = "gold",
+    stroke_color: str = "#666",
+    stroke_width_px: int = 1,
+) -> str:
+    direction = "LR" if horizontal else "TD"
+    lines: list[str] = [f"flowchart {direction}"]
+    if title:
+        lines.append(f"%% {title}")
+
+    result_id: int = id(tensor)
+    visited: set[int] = set()
+    nodes_to_draw: list[Tensor] = []
+
+    def dfs(t: Tensor) -> None:
+        if id(t) in visited:
+            return
+        visited.add(id(t))
+        for p in t._prev:
+            dfs(p)
+        nodes_to_draw.append(t)
+
+    def tensor_node_id(t: Tensor) -> str:
+        return f"t_{id(t)}"
+
+    def op_node_id(op: object) -> str:
+        return f"op_{id(op)}"
+
+    def add_node(node_id: str, label: str, kind: str) -> None:
+        if node_id in defined_nodes:
+            return
+        defined_nodes.add(node_id)
+        if kind == "op":
+            lines.append(f'{node_id}(("{label}"))')
+        else:
+            lines.append(f'{node_id}["{label}"]')
+
+    dfs(tensor)
+
+    defined_nodes: set[str] = set()
+    edge_lines: list[str] = []
+    class_lines: list[str] = []
+
+    for t in nodes_to_draw:
+        t_id = tensor_node_id(t)
+
+        if not t.is_leaf and t._op is not None:
+            op_id = op_node_id(t._op)
+            op_label = type(t._op).__name__
+            add_node(op_id, op_label, "op")
+            edge_lines.append(f"{op_id} --> {t_id}")
+            class_lines.append(f"class {op_id} op")
+            for inp in t._prev:
+                edge_lines.append(f"{tensor_node_id(inp)} --> {op_id}")
+
+        shape_label = str(t.shape) if t.ndim > 0 else str(t.item())
+        add_node(t_id, shape_label, "tensor")
+
+        if start_id is not None and id(t) == start_id:
+            class_lines.append(f"class {t_id} start")
+        elif isinstance(t, nn.Parameter):
+            class_lines.append(f"class {t_id} param")
+        elif id(t) == result_id:
+            class_lines.append(f"class {t_id} result")
+        elif not t.requires_grad:
+            class_lines.append(f"class {t_id} leaf")
+        else:
+            class_lines.append(f"class {t_id} grad")
+
+    lines.extend(edge_lines)
+    if use_class_defs:
+        lines.append(
+            f"classDef op fill:{op_fill},stroke:{stroke_color},stroke-width:{stroke_width_px}px;"
+        )
+        lines.append(
+            f"classDef param fill:{param_fill},stroke:{stroke_color},stroke-width:{stroke_width_px}px;"
+        )
+        lines.append(
+            f"classDef result fill:{result_fill},stroke:{stroke_color},stroke-width:{stroke_width_px}px;"
+        )
+        lines.append(
+            f"classDef leaf fill:{leaf_fill},stroke:{stroke_color},stroke-width:{stroke_width_px}px;"
+        )
+        lines.append(
+            f"classDef grad fill:{grad_fill},stroke:{stroke_color},stroke-width:{stroke_width_px}px;"
+        )
+        lines.append(
+            f"classDef start fill:{start_fill},stroke:{stroke_color},stroke-width:{stroke_width_px}px;"
+        )
+        lines.extend(class_lines)
+
+    if end_semicolons:
+        lines = [
+            f"{line};" if line and not line.endswith(";") else line for line in lines
+        ]
+
+    text = "\n".join(lines)
+    if copy_to_clipboard:
+        _copy_to_clipboard(text)
+    return text
+
+
 def _copy_to_clipboard(text: str) -> None:
     import os
     import shutil

{lucid_dl-2.11.3 → lucid_dl-2.11.5/lucid_dl.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lucid-dl
-Version: 2.11.3
+Version: 2.11.5
 Summary: Lumerico's Comprehensive Interface for Deep Learning
 Home-page: https://github.com/ChanLumerico/lucid
 Author: ChanLumerico
@@ -48,19 +48,9 @@ Whether you're a student, educator, or an advanced researcher seeking to demysti
 
 ### 🔥 What's New
 
-- Added additional `nn.Module` hooks for richer introspection during training:
-
-  ```python
-  def register_forward_pre_hook(self, hook: Callable, *, with_kwargs: bool = False)
-
-  def register_forward_hook(self, hook: Callable, *, with_kwargs: bool = False)
+- Added new visual tool: `lucid.visual.build_tensor_mermaid_chart` which builds a Mermaid chart of given tensor's computatoinal graph
 
-  def register_backward_hook(self, hook: Callable)
-
-  def register_full_backward_pre_hook(self, hook: Callable)
-
-  def register_full_backward_hook(self, hook: Callable)
-  ```
+- Added additional `nn.Module` hooks for richer introspection during training:
 
 ## 🔧 How to Install
 

{lucid_dl-2.11.3 → lucid_dl-2.11.5}/lucid_dl.egg-info/SOURCES.txt

@@ -134,7 +134,6 @@ lucid/transforms/__init__.py
 lucid/transforms/_base.py
 lucid/transforms/image.py
 lucid/visual/__init__.py
-lucid/visual/graph.py
 lucid/visual/mermaid.py
 lucid/weights/__init__.py
 lucid/weights/__init__.pyi

{lucid_dl-2.11.3 → lucid_dl-2.11.5}/setup.py

@@ -5,7 +5,7 @@ with open("README.md", "r") as fh:
 
 setuptools.setup(
     name="lucid-dl",
-    version="2.11.3",
+    version="2.11.5",
     author="ChanLumerico",
     author_email="greensox284@gmail.com",
     description="Lumerico's Comprehensive Interface for Deep Learning",

lucid_dl-2.11.3/lucid/visual/graph.py

@@ -1,141 +0,0 @@
-from typing import Union
-from warnings import deprecated
-
-import networkx as nx
-import matplotlib.pyplot as plt
-
-import lucid.nn as nn
-from lucid._tensor import Tensor
-
-
-__all__ = ["draw_tensor_graph"]
-
-
-@deprecated("This feature will be re-written with Mermaid in future relases.")
-def draw_tensor_graph(
-    tensor: Tensor,
-    horizontal: bool = False,
-    title: Union[str, None] = None,
-    start_id: Union[int, None] = None,
-) -> plt.Figure:
-    G: nx.DiGraph = nx.DiGraph()
-    result_id: int = id(tensor)
-
-    visited: set[int] = set()
-    nodes_to_draw: list[Tensor] = []
-
-    def dfs(t: Tensor) -> None:
-        if id(t) in visited:
-            return
-        visited.add(id(t))
-        for p in t._prev:
-            dfs(p)
-        nodes_to_draw.append(t)
-
-    dfs(tensor)
-
-    for t in nodes_to_draw:
-        if not t.is_leaf and t._op is not None:
-            op_id: int = id(t._op)
-            op_label: str = type(t._op).__name__
-            G.add_node(op_id, label=op_label, shape="circle", color="lightgreen")
-            G.add_edge(op_id, id(t))
-            for inp in t._prev:
-                G.add_edge(id(inp), op_id)
-
-        shape_label: str = str(t.shape) if t.ndim > 0 else str(t.item())
-        if isinstance(t, nn.Parameter):
-            color: str = "plum"
-        else:
-            color = (
-                "lightcoral"
-                if id(t) == result_id
-                else "lightgray" if not t.requires_grad else "lightblue"
-            )
-        if start_id is not None and id(t) == start_id:
-            color = "gold"
-
-        G.add_node(id(t), label=shape_label, shape="rectangle", color=color)
-
-    def grid_layout(
-        G: nx.DiGraph, horizontal: bool = False
-    ) -> tuple[dict, tuple, float, int]:
-        levels: dict[int, int] = {}
-        for node in nx.topological_sort(G):
-            preds = list(G.predecessors(node))
-            levels[node] = 0 if not preds else max(levels[p] for p in preds) + 1
-
-        level_nodes: dict[int, list[int]] = {}
-        for node, level in levels.items():
-            level_nodes.setdefault(level, []).append(node)
-
-        def autoscale(
-            level_nodes: dict[int, list[int]],
-            horizontal: bool = False,
-            base_size: float = 0.5,
-            base_nodesize: int = 500,
-        ) -> tuple[tuple[float, float], float, int]:
-            num_levels: int = len(level_nodes)
-            max_width: int = max(len(nodes) for nodes in level_nodes.values())
-            node_count: int = sum(len(nodes) for nodes in level_nodes.values())
-
-            if horizontal:
-                fig_w: float = min(32, max(4.0, base_size * num_levels))
-                fig_h: float = min(32, max(4.0, base_size * max_width))
-            else:
-                fig_w = min(32, max(4.0, base_size * max_width))
-                fig_h = min(32, max(4.0, base_size * num_levels))
-
-            nodesize: float = (
-                base_nodesize
-                if node_count <= 100
-                else base_nodesize * (100 / node_count)
-            )
-            fontsize: int = max(5, min(8, int(80 / node_count)))
-            return (fig_w, fig_h), nodesize, fontsize
-
-        figsize, nodesize, fontsize = autoscale(level_nodes, horizontal)
-        pos: dict[int, tuple[float, float]] = {}
-        for level, nodes in level_nodes.items():
-            for i, node in enumerate(nodes):
-                pos[node] = (
-                    (level * 2.5, -i * 2.0) if horizontal else (i * 2.5, -level * 2.0)
-                )
-        return pos, figsize, nodesize, fontsize
-
-    labels: dict[int, str] = nx.get_node_attributes(G, "label")
-    colors: dict[int, str] = nx.get_node_attributes(G, "color")
-    shapes: dict[int, str] = nx.get_node_attributes(G, "shape")
-    pos, figsize, nodesize, fontsize = grid_layout(G, horizontal)
-
-    fig, ax = plt.subplots(figsize=figsize)
-
-    rect_nodes: list[int] = [n for n in G.nodes() if shapes.get(n) == "rectangle"]
-    circ_nodes: list[int] = [n for n in G.nodes() if shapes.get(n) == "circle"]
-    rect_colors: list[str] = [colors[n] for n in rect_nodes]
-
-    nx.draw_networkx_nodes(
-        G,
-        pos,
-        nodelist=rect_nodes,
-        node_color=rect_colors,
-        node_size=nodesize,
-        node_shape="s",
-        ax=ax,
-    )
-    nx.draw_networkx_nodes(
-        G,
-        pos,
-        nodelist=circ_nodes,
-        node_color="lightgreen",
-        node_size=nodesize,
-        node_shape="o",
-        ax=ax,
-    )
-    nx.draw_networkx_edges(G, pos, width=0.5, arrows=True, edge_color="gray", ax=ax)
-    nx.draw_networkx_labels(G, pos, labels=labels, font_size=fontsize, ax=ax)
-
-    ax.axis("off")
-    ax.set_title(title if title is not None else "")
-
-    return fig