Skip to content

utorch.utils

utorch.utils

__all__ module-attribute 

__all__ = ['count_nodes', 'draw_graph', 'show_drawing_canvas']

count_nodes 

count_nodes(root)

Traverse the computational DAG and count leaf and interim nodes.

Parameters:

Name Type Description Default
root Tensor

The output node from which to trace the graph backwards.

 required

Returns:

Type Description
GraphStats

GraphStats containing leaf count, interim count, edge count, and op breakdown.
Example

x = Tensor(2.0, name="x") y = Tensor(3.0, name="y") z = x * y + 1.0 stats = count_nodes(z) stats.leaves 3 # x, y, and the constant 1.0 stats.interim 2 # the '*' node and the '+' node

Source code in utorch/utils/graph_stats.py 
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
def count_nodes(root: Tensor) -> GraphStats:
    """Traverse the computational DAG and count leaf and interim nodes.

    Args:
        root: The output node from which to trace the graph backwards.

    Returns:
        GraphStats containing leaf count, interim count, edge count, and op breakdown.

    Example:
        >>> x = Tensor(2.0, name="x")
        >>> y = Tensor(3.0, name="y")
        >>> z = x * y + 1.0
        >>> stats = count_nodes(z)
        >>> stats.leaves
        3  # x, y, and the constant 1.0
        >>> stats.interim
        2  # the '*' node and the '+' node
    """
    seen: set[int] = set()
    leaf_count = 0
    interim_count = 0
    edge_count = 0
    ops: dict[str, int] = {}

    def traverse(node: Tensor) -> None:
        nonlocal leaf_count, interim_count, edge_count

        node_id = id(node)
        if node_id in seen:
            return
        seen.add(node_id)

        if node.is_leaf:
            leaf_count += 1
        else:
            interim_count += 1
            op_name = getattr(node, "op", None) or "func"
            ops[op_name] = ops.get(op_name, 0) + 1

        for parent in node.parents:
            edge_count += 1
            traverse(parent)

    traverse(root)

    return GraphStats(
        leaves=leaf_count,
        interim=interim_count,
        total=leaf_count + interim_count,
        edges=edge_count,
        ops=ops,
    )

draw_graph 

draw_graph(root, format='png', rankdir='LR')

Render a computational graph with a tree-inspired aesthetic.

  • Leaf nodes (inputs, constants, parameters) are rendered in green.
  • Interim nodes (operations, intermediate results) are rendered in brown.

Parameters:

Name Type Description Default
root Tensor

The output node (root of the evaluation) to trace backwards.

 required
format str

Output rendering format ('svg', 'png', 'pdf').

 'png'
rankdir str

Graph layout direction ('LR' = left-to-right, 'TB' = top-to-bottom).

 'LR'

Returns:

Type Description
Any

graphviz.Digraph object ready for notebook rendering or file saving.
Source code in utorch/utils/vizualization.py 
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
def draw_graph(
    root: Tensor,
    format: str = "png",
    rankdir: str = "LR",
) -> Any:
    """Render a computational graph with a tree-inspired aesthetic.

    - Leaf nodes (inputs, constants, parameters) are rendered in green.
    - Interim nodes (operations, intermediate results) are rendered in brown.

    Args:
        root: The output node (root of the evaluation) to trace backwards.
        format: Output rendering format ('svg', 'png', 'pdf').
        rankdir: Graph layout direction ('LR' = left-to-right, 'TB' = top-to-bottom).

    Returns:
        graphviz.Digraph object ready for notebook rendering or file saving.
    """

    # Graph-level attributes
    dot = Digraph(
        name="computation_graph",
        format=format,
        graph_attr={
            "rankdir": rankdir,
            "splines": "ortho",
            "nodesep": "0.4",
            "ranksep": "0.6",
            "bgcolor": "transparent",
        },
        node_attr={
            "style": "filled,rounded",
            "shape": "box",
            "fontname": "Helvetica, Arial, sans-serif",
            "fontsize": "11",
            "penwidth": "1.5",
        },
        edge_attr={
            "color": "#5D4037",
            "arrowsize": "0.8",
            "penwidth": "1.2",
        },
    )

    seen: set[int] = set()

    def format_value(val: Any) -> str:
        """Format node values cleanly for diagram labels."""
        if isinstance(val, np.ndarray):
            if val.size == 1:
                return f"{val.item():.4g}"
            return f"array(shape={val.shape})"
        if isinstance(val, int | float):
            return f"{val:.4g}"
        return str(val)

    def build(node: Tensor) -> None:
        node_id = str(id(node))
        if id(node) in seen:
            return
        seen.add(id(node))

        val_str = format_value(node.value)

        if node.is_leaf:
            # --- LEAF NODE (Green) ---
            title = f"{node.name}\n" if node.name else ""
            label = f"{title}data: {val_str}"
            dot.node(
                node_id,
                label=label,
                fillcolor="#C8E6C9",  # Soft leaf green
                color="#2E7D32",  # Forest green border
                fontcolor="#1B5E20",  # Dark green text
            )
        else:
            # --- INTERIM NODE (Woody Brown) ---
            op_name = _get_clean_op_name(node)
            op_str = f"fun: {op_name}" if node.fun else "op: func"
            name_str = f"{node.name} | " if node.name else ""
            label = f"{name_str}{op_str}\nval: {val_str}"
            dot.node(
                node_id,
                label=label,
                fillcolor="#D7CCC8",  # Light wood brown
                color="#5D4037",  # Bark brown border
                fontcolor="#3E2723",  # Deep espresso text
            )

        # Draw directed edges from parents -> current node
        for parent in node.parents:
            parent_id = str(id(parent))
            dot.edge(parent_id, node_id)
            build(parent)

    build(root)
    return dot

show_drawing_canvas 

show_drawing_canvas()

Render an interactive 28x28 drawing canvas.

Source code in utorch/utils/notebook_utils.py 
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
def show_drawing_canvas() -> None:
    """Render an interactive 28x28 drawing canvas."""
    html_code = """
    <div style="background: #1e1e1e; padding: 15px; border-radius: 8px; display: inline-block; font-family: sans-serif;">
      <h4 style="color: #eee; margin: 0 0 10px 0;">Draw a digit (0-9) with your mouse:</h4>
      <canvas id="cvs" width="280" height="280" style="background: black; border: 2px solid #ee4c2c; cursor: crosshair; border-radius: 4px;"></canvas>
      <div style="margin-top: 10px;">
        <button onclick="clearCanvas()" style="padding: 6px 14px; font-weight: bold; background: #444; color: white; border: none; border-radius: 4px; cursor: pointer;">Clear</button>
        <button onclick="captureCanvas()" style="padding: 6px 14px; font-weight: bold; background: #ee4c2c; color: white; border: none; border-radius: 4px; cursor: pointer; margin-left: 8px;">1. Process Image</button>
      </div>
      <div id="output_box" style="margin-top: 12px; display: none;">
        <span style="color: #aaa; font-size: 12px;">Click below to copy pixels to clipboard, then paste in the next cell:</span><br/>
        <textarea id="pixel_txt" rows="3" cols="35" style="width: 270px; background: #111; color: #4af; font-size: 11px; border: 1px solid #333; margin-top: 4px;"></textarea>
        <br/>
        <button onclick="copyToClipboard()" style="margin-top: 6px; padding: 5px 12px; font-size: 12px; background: #2ca02c; color: white; border: none; border-radius: 4px; cursor: pointer;">Copy Code</button>
      </div>
    </div>

    <script>
      (function() {
        const c = document.getElementById('cvs');
        const ctx = c.getContext('2d');
        let drawing = false;

        ctx.strokeStyle = "white";
        ctx.lineWidth = 22;
        ctx.lineCap = "round";
        ctx.lineJoin = "round";

        function getPos(e) {
          const r = c.getBoundingClientRect();
          return { x: e.clientX - r.left, y: e.clientY - r.top };
        }

        c.onmousedown = (e) => { drawing = true; const p = getPos(e); ctx.beginPath(); ctx.moveTo(p.x, p.y); };
        window.onmouseup = () => { drawing = false; };
        c.onmousemove = (e) => {
          if (!drawing) return;
          const p = getPos(e);
          ctx.lineTo(p.x, p.y);
          ctx.stroke();
        };

        window.clearCanvas = function() {
          ctx.fillStyle = "black";
          ctx.fillRect(0, 0, 280, 280);
          document.getElementById('output_box').style.display = "none";
        };
        clearCanvas();

        window.captureCanvas = function() {
          // Downscale to 28x28
          const s = document.createElement('canvas');
          s.width = 28; s.height = 28;
          const sctx = s.getContext('2d');
          sctx.drawImage(c, 0, 0, 28, 28);

          const img = sctx.getImageData(0, 0, 28, 28).data;
          const vals = [];
          for (let i = 0; i < img.length; i += 4) {
            vals.push(Math.round((img[i] / 255.0) * 100) / 100);
          }

          const code = "drawn_digit = np.array(" + JSON.stringify(vals) + ").reshape(1, 784)";
          const box = document.getElementById('output_box');
          const txt = document.getElementById('pixel_txt');
          txt.value = code;
          box.style.display = "block";
        };

        window.copyToClipboard = function() {
          const txt = document.getElementById('pixel_txt');
          txt.select();
          document.execCommand('copy');
          alert("Copied code to clipboard! Paste and run it in the next cell.");
        };
      })();
    </script>
    """
    display(HTML(html_code))  # type: ignore[no-untyped-call]

data

DataLoader 

DataLoader(dataset, batch_size=1, shuffle=False, drop_last=False)

Data loader combining a dataset and sampler, yielding batched Tensors.

Parameters:

Name Type Description Default
dataset Dataset

Dataset from which to load data.

 required
batch_size int

How many samples per batch to load (default: 1).

 1
shuffle bool

If True, reshuffles indices at the start of every epoch.

 False
drop_last bool

If True, drops the last incomplete batch.

 False

Initialize a DataLoader.

Source code in utorch/utils/data.py 
64
65
66
67
68
69
70
71
72
73
74
75
def __init__(
    self,
    dataset: Dataset,
    batch_size: int = 1,
    shuffle: bool = False,
    drop_last: bool = False,
) -> None:
    """Initialize a DataLoader."""
    self.dataset = dataset
    self.batch_size = batch_size
    self.shuffle = shuffle
    self.drop_last = drop_last
__iter__ 
__iter__()

Return an iterator over the samples in the dataset.

Source code in utorch/utils/data.py 
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
def __iter__(self) -> Iterator[tuple[Tensor, ...]]:
    """Return an iterator over the samples in the dataset."""
    indices = np.arange(len(self.dataset))
    if self.shuffle:
        np.random.shuffle(indices)

    for start_idx in range(0, len(self.dataset), self.batch_size):
        batch_indices = indices[start_idx : start_idx + self.batch_size]
        if self.drop_last and len(batch_indices) < self.batch_size:
            continue

        samples = [self.dataset[i] for i in batch_indices]
        num_fields = len(samples[0])
        batch = tuple(
            Tensor(np.stack([s[f].value for s in samples], axis=0))
            for f in range(num_fields)
        )
        yield batch
__len__ 
__len__()

Return the total number of samples in the dataset.

Source code in utorch/utils/data.py 
77
78
79
80
81
def __len__(self) -> int:
    """Return the total number of samples in the dataset."""
    if self.drop_last:
        return len(self.dataset) // self.batch_size
    return (len(self.dataset) + self.batch_size - 1) // self.batch_size

Dataset

An abstract class representing a dataset.

All custom datasets must override __getitem__ (returning a sample at an index) and __len__ (returning the total size of the dataset).

__getitem__ 
__getitem__(index)

Override this method to return a sample at an index.

Source code in utorch/utils/data.py 
18
19
20
def __getitem__(self, index: int) -> Any:
    """Override this method to return a sample at an index."""
    raise NotImplementedError
__len__ 
__len__()

Override this method to return the total size of the dataset.

Source code in utorch/utils/data.py 
22
23
24
def __len__(self) -> int:
    """Override this method to return the total size of the dataset."""
    raise NotImplementedError

TensorDataset 

TensorDataset(*tensors)

Bases: Dataset

Dataset wrapping tensors where indexing retrieves slices along axis 0.

Parameters:

Name Type Description Default
*tensors Any

Tensors or NumPy arrays with identical length along axis 0.

 ()

Initialize a TensorDataset.

Source code in utorch/utils/data.py 
34
35
36
37
38
39
40
41
42
43
def __init__(self, *tensors: Any) -> None:
    """Initialize a TensorDataset."""
    assert len(tensors) > 0, "TensorDataset requires at least one tensor."
    first_len = len(tensors[0])
    for t in tensors:
        assert len(t) == first_len, (
            "All tensors must have identical size along dim 0."
        )

    self.tensors = [t if isinstance(t, Tensor) else Tensor(t) for t in tensors]
__getitem__ 
__getitem__(index)

Return a single tensor example from the dataset.

Source code in utorch/utils/data.py 
45
46
47
def __getitem__(self, index: int) -> tuple[Tensor, ...]:
    """Return a single tensor example from the dataset."""
    return tuple(Tensor(t.value[index]) for t in self.tensors)
__len__ 
__len__()

Return the number of tensors in the dataset.

Source code in utorch/utils/data.py 
49
50
51
def __len__(self) -> int:
    """Return the number of tensors in the dataset."""
    return len(self.tensors[0].value)

graph_stats

GraphStats dataclass 

GraphStats(leaves, interim, total, edges, ops)

Summary statistics for a computational graph.

Attributes:

Name Type Description
leaves int

Number of input/leaf nodes (parameters, inputs, constants).
interim int

Number of operation/intermediate nodes.
total int

Total unique nodes in the graph (leaves + interim).
edges int

Total dependency connections (edges) between nodes.
ops dict[str, int]

Breakdown of operations by frequency (e.g. {'add': 2, 'mul': 1}).
__repr__ 
__repr__()

Return a string representation of the Graph stat for inspection.

Source code in utorch/utils/graph_stats.py 
26
27
28
29
30
31
32
def __repr__(self) -> str:
    """Return a string representation of the Graph stat for inspection."""
    ops_str = ", ".join(f"{op}: {count}" for op, count in self.ops.items())
    return (
        f"GraphStats(total={self.total}, leaves={self.leaves}, "
        f"interim={self.interim}, edges={self.edges}, ops={{{ops_str}}})"
    )

count_nodes 

count_nodes(root)

Traverse the computational DAG and count leaf and interim nodes.

Parameters:

Name Type Description Default
root Tensor

The output node from which to trace the graph backwards.

 required

Returns:

Type Description
GraphStats

GraphStats containing leaf count, interim count, edge count, and op breakdown.
Example

x = Tensor(2.0, name="x") y = Tensor(3.0, name="y") z = x * y + 1.0 stats = count_nodes(z) stats.leaves 3 # x, y, and the constant 1.0 stats.interim 2 # the '*' node and the '+' node

Source code in utorch/utils/graph_stats.py 
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
def count_nodes(root: Tensor) -> GraphStats:
    """Traverse the computational DAG and count leaf and interim nodes.

    Args:
        root: The output node from which to trace the graph backwards.

    Returns:
        GraphStats containing leaf count, interim count, edge count, and op breakdown.

    Example:
        >>> x = Tensor(2.0, name="x")
        >>> y = Tensor(3.0, name="y")
        >>> z = x * y + 1.0
        >>> stats = count_nodes(z)
        >>> stats.leaves
        3  # x, y, and the constant 1.0
        >>> stats.interim
        2  # the '*' node and the '+' node
    """
    seen: set[int] = set()
    leaf_count = 0
    interim_count = 0
    edge_count = 0
    ops: dict[str, int] = {}

    def traverse(node: Tensor) -> None:
        nonlocal leaf_count, interim_count, edge_count

        node_id = id(node)
        if node_id in seen:
            return
        seen.add(node_id)

        if node.is_leaf:
            leaf_count += 1
        else:
            interim_count += 1
            op_name = getattr(node, "op", None) or "func"
            ops[op_name] = ops.get(op_name, 0) + 1

        for parent in node.parents:
            edge_count += 1
            traverse(parent)

    traverse(root)

    return GraphStats(
        leaves=leaf_count,
        interim=interim_count,
        total=leaf_count + interim_count,
        edges=edge_count,
        ops=ops,
    )

notebook_utils

Interactive 28x28 drawing canvas for VS Code, JupyterLab, and Colab.

show_drawing_canvas 

show_drawing_canvas()

Render an interactive 28x28 drawing canvas.

Source code in utorch/utils/notebook_utils.py 
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
def show_drawing_canvas() -> None:
    """Render an interactive 28x28 drawing canvas."""
    html_code = """
    <div style="background: #1e1e1e; padding: 15px; border-radius: 8px; display: inline-block; font-family: sans-serif;">
      <h4 style="color: #eee; margin: 0 0 10px 0;">Draw a digit (0-9) with your mouse:</h4>
      <canvas id="cvs" width="280" height="280" style="background: black; border: 2px solid #ee4c2c; cursor: crosshair; border-radius: 4px;"></canvas>
      <div style="margin-top: 10px;">
        <button onclick="clearCanvas()" style="padding: 6px 14px; font-weight: bold; background: #444; color: white; border: none; border-radius: 4px; cursor: pointer;">Clear</button>
        <button onclick="captureCanvas()" style="padding: 6px 14px; font-weight: bold; background: #ee4c2c; color: white; border: none; border-radius: 4px; cursor: pointer; margin-left: 8px;">1. Process Image</button>
      </div>
      <div id="output_box" style="margin-top: 12px; display: none;">
        <span style="color: #aaa; font-size: 12px;">Click below to copy pixels to clipboard, then paste in the next cell:</span><br/>
        <textarea id="pixel_txt" rows="3" cols="35" style="width: 270px; background: #111; color: #4af; font-size: 11px; border: 1px solid #333; margin-top: 4px;"></textarea>
        <br/>
        <button onclick="copyToClipboard()" style="margin-top: 6px; padding: 5px 12px; font-size: 12px; background: #2ca02c; color: white; border: none; border-radius: 4px; cursor: pointer;">Copy Code</button>
      </div>
    </div>

    <script>
      (function() {
        const c = document.getElementById('cvs');
        const ctx = c.getContext('2d');
        let drawing = false;

        ctx.strokeStyle = "white";
        ctx.lineWidth = 22;
        ctx.lineCap = "round";
        ctx.lineJoin = "round";

        function getPos(e) {
          const r = c.getBoundingClientRect();
          return { x: e.clientX - r.left, y: e.clientY - r.top };
        }

        c.onmousedown = (e) => { drawing = true; const p = getPos(e); ctx.beginPath(); ctx.moveTo(p.x, p.y); };
        window.onmouseup = () => { drawing = false; };
        c.onmousemove = (e) => {
          if (!drawing) return;
          const p = getPos(e);
          ctx.lineTo(p.x, p.y);
          ctx.stroke();
        };

        window.clearCanvas = function() {
          ctx.fillStyle = "black";
          ctx.fillRect(0, 0, 280, 280);
          document.getElementById('output_box').style.display = "none";
        };
        clearCanvas();

        window.captureCanvas = function() {
          // Downscale to 28x28
          const s = document.createElement('canvas');
          s.width = 28; s.height = 28;
          const sctx = s.getContext('2d');
          sctx.drawImage(c, 0, 0, 28, 28);

          const img = sctx.getImageData(0, 0, 28, 28).data;
          const vals = [];
          for (let i = 0; i < img.length; i += 4) {
            vals.push(Math.round((img[i] / 255.0) * 100) / 100);
          }

          const code = "drawn_digit = np.array(" + JSON.stringify(vals) + ").reshape(1, 784)";
          const box = document.getElementById('output_box');
          const txt = document.getElementById('pixel_txt');
          txt.value = code;
          box.style.display = "block";
        };

        window.copyToClipboard = function() {
          const txt = document.getElementById('pixel_txt');
          txt.select();
          document.execCommand('copy');
          alert("Copied code to clipboard! Paste and run it in the next cell.");
        };
      })();
    </script>
    """
    display(HTML(html_code))  # type: ignore[no-untyped-call]

vizualization

Visualization utilities for computational graphs.

draw_graph 

draw_graph(root, format='png', rankdir='LR')

Render a computational graph with a tree-inspired aesthetic.

  • Leaf nodes (inputs, constants, parameters) are rendered in green.
  • Interim nodes (operations, intermediate results) are rendered in brown.

Parameters:

Name Type Description Default
root Tensor

The output node (root of the evaluation) to trace backwards.

 required
format str

Output rendering format ('svg', 'png', 'pdf').

 'png'
rankdir str

Graph layout direction ('LR' = left-to-right, 'TB' = top-to-bottom).

 'LR'

Returns:

Type Description
Any

graphviz.Digraph object ready for notebook rendering or file saving.
Source code in utorch/utils/vizualization.py 
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
def draw_graph(
    root: Tensor,
    format: str = "png",
    rankdir: str = "LR",
) -> Any:
    """Render a computational graph with a tree-inspired aesthetic.

    - Leaf nodes (inputs, constants, parameters) are rendered in green.
    - Interim nodes (operations, intermediate results) are rendered in brown.

    Args:
        root: The output node (root of the evaluation) to trace backwards.
        format: Output rendering format ('svg', 'png', 'pdf').
        rankdir: Graph layout direction ('LR' = left-to-right, 'TB' = top-to-bottom).

    Returns:
        graphviz.Digraph object ready for notebook rendering or file saving.
    """

    # Graph-level attributes
    dot = Digraph(
        name="computation_graph",
        format=format,
        graph_attr={
            "rankdir": rankdir,
            "splines": "ortho",
            "nodesep": "0.4",
            "ranksep": "0.6",
            "bgcolor": "transparent",
        },
        node_attr={
            "style": "filled,rounded",
            "shape": "box",
            "fontname": "Helvetica, Arial, sans-serif",
            "fontsize": "11",
            "penwidth": "1.5",
        },
        edge_attr={
            "color": "#5D4037",
            "arrowsize": "0.8",
            "penwidth": "1.2",
        },
    )

    seen: set[int] = set()

    def format_value(val: Any) -> str:
        """Format node values cleanly for diagram labels."""
        if isinstance(val, np.ndarray):
            if val.size == 1:
                return f"{val.item():.4g}"
            return f"array(shape={val.shape})"
        if isinstance(val, int | float):
            return f"{val:.4g}"
        return str(val)

    def build(node: Tensor) -> None:
        node_id = str(id(node))
        if id(node) in seen:
            return
        seen.add(id(node))

        val_str = format_value(node.value)

        if node.is_leaf:
            # --- LEAF NODE (Green) ---
            title = f"{node.name}\n" if node.name else ""
            label = f"{title}data: {val_str}"
            dot.node(
                node_id,
                label=label,
                fillcolor="#C8E6C9",  # Soft leaf green
                color="#2E7D32",  # Forest green border
                fontcolor="#1B5E20",  # Dark green text
            )
        else:
            # --- INTERIM NODE (Woody Brown) ---
            op_name = _get_clean_op_name(node)
            op_str = f"fun: {op_name}" if node.fun else "op: func"
            name_str = f"{node.name} | " if node.name else ""
            label = f"{name_str}{op_str}\nval: {val_str}"
            dot.node(
                node_id,
                label=label,
                fillcolor="#D7CCC8",  # Light wood brown
                color="#5D4037",  # Bark brown border
                fontcolor="#3E2723",  # Deep espresso text
            )

        # Draw directed edges from parents -> current node
        for parent in node.parents:
            parent_id = str(id(parent))
            dot.edge(parent_id, node_id)
            build(parent)

    build(root)
    return dot