code: def _graph_to_dot(graph: Graph, args, outputs) -> str:
      """Converts from an internal graph IR to 'dot' format."""
      if tree is None:
          raise ImportError("hk.to_dot requires dm-tree>=0.1.1.")
  
      def format_path(path):
          if isinstance(outputs, tuple):
              out = f"output[{path[0]}]"
              if len(path) > 1:
                  out += ": " + "/".join(map(str, path[1:]))
          else:
              out = "output"
              if path:
                  out += ": " + "/".join(map(str, path))
          return out
  
      lines = []
      used_argids = set()
      argid_usecount = collections.Counter()
      op_outids = set()
      captures = []
      argids = {id(v) for v in jax.tree_util.tree_leaves(args)}
      outids = {id(v) for v in jax.tree_util.tree_leaves(outputs)}
      outname = {
          id(v): format_path(p)
          for p, v in tree.flatten_with_path(outputs)
      }
  
      def render_graph(
          g: Graph, parent: Optional[Graph] = None, depth: int = 0
      ):
          """Renders a given graph by appending 'dot' format lines."""
  
          if parent:
              lines.extend(
                  [
                      f"subgraph cluster_{id(g)} {{",
                      '  style="rounded,filled";',
                      '  fillcolor="#F0F5F5";',
                      '  color="#14234B;";',
                      "  pad=0.1;",
                      f"  fontsize={_scaled_font_size(depth)};",
                      f"  label = <<b>{escape(g.title)}</b>>;",
                      "  labelloc = t;",
                  ]
              )
  
          for node in g.nodes:
              label = f"<b>{escape(node.title)}</b>"
              for o in node.outputs:
                  label += "<br/>" + _format_val(o)
                  op_outids.add(id(o))
  
              node_id = id(node.id)
              if node_id in outids:
                  label = (
                      f"<b>{escape(outname[node_id])}</b><br/>" + label
                  )
                  color = "#0053D6"
                  fillcolor = "#AABFFF"
                  style = "filled,bold"
              else:
                  color = "#FFDB13"
                  fillcolor = "#FFF26E"
                  style = "filled"
  
              lines.append(
                  f"{node_id} [label=<{label}>, "
                  f' id="node{node_id}",'
                  " shape=rect,"
                  f' style="{style}",'
                  ' tooltip=" ",'
                  ' fontcolor="black",'
                  f' color="{color}",'
                  f' fillcolor="{fillcolor}"];'
              )
  
          for s in g.subgraphs:
              render_graph(s, parent=g, depth=depth - 1)
  
          if parent:
              lines.append(f"}}  // subgraph cluster_{id(g)}")
  
          for a, b in g.edges:
              if id(a) not in argids and id(a) not in op_outids:
                  captures.append(a)
  
              a, b = map(id, (a, b))
              if a in argids:
                  i = argid_usecount[a]
                  argid_usecount[a] += 1
                  lines.append(f"{a}{i} -> {b};")
              else:
                  lines.append(f"{a} -> {b};")
              used_argids.add(a)
  
      graph_depth = _max_depth(graph)
      render_graph(graph, parent=None, depth=graph_depth)
  
      # Process inputs and label them in the graph.
      for path, value in tree.flatten_with_path(args):
          if value is None:
              continue
  
          node_id = id(value)
          if node_id not in used_argids:
              continue
  
          for i in range(argid_usecount[node_id]):
              label = f"<b>args[{escape(path[0])}]"
              if len(path) > 1:
                  label += ": " + "/".join(map(str, path[1:]))
              label += "</b>"
              if hasattr(value, "shape") and hasattr(value, "dtype"):
                  label += f"<br/>{escape(_format_val(value))}"
              fillcolor = "#FFDEAF"
              fontcolor = "black"
  
              if i > 0:
                  label = "<b>(reuse)</b><br/>" + label
                  fillcolor = "#FFEACC"
                  fontcolor = "#565858"
  
              lines.append(
                  f"{node_id}{i} [label=<{label}>"
                  f' id="node{node_id}{i}",'
                  " shape=rect,"
                  ' style="filled",'
                  f' fontcolor="{fontcolor}",'
                  ' color="#FF8A4F",'
                  f' fillcolor="{fillcolor}"];'
              )
  
      for value in captures:
          node_id = id(value)
          if (
              not hasattr(value, "aval")
              and hasattr(value, "size")
              and value.size == 1
          ):
              label = f"<b>{value.item()}</b>"
          else:
              label = f"<b>{escape(_format_val(value))}</b>"
  
          lines.append(
              f"{node_id} [label=<{label}>"
              " shape=rect,"
              ' style="filled",'
              ' fontcolor="black",'
              ' color="#A261FF",'
              ' fillcolor="#E6D6FF"];'
          )
  
      head = [
          "digraph G {",
          "rankdir = TD;",
          "compound = true;",
          f"label = <<b>{escape(graph.title)}</b>>;",
          f"fontsize={_scaled_font_size(graph_depth)};",
          "labelloc = t;",
          "stylesheet = <",
          "  data:text/css,",
          "  @import url(https://fonts.googleapis.com/css?family=Roboto:400,700);",
          "  svg text {",
          "    font-family: 'Roboto';",
          "  }",
          "  .node text {",
          "    font-size: 12px;",
          "  }",
      ]
      for node_id, use_count in argid_usecount.items():
          if use_count == 1:
              continue
          # Add hover animation for reused args.
          for a in range(use_count):
              for b in range(use_count):
                  if a == b:
                      head.append(
                          f"%23node{node_id}{a}:hover "
                          "{ stroke-width: 0.2em; }"
                      )
                  else:
                      head.append(
                          f"%23node{node_id}{a}:hover ~ %23node{node_id}{b} "
                          "{ stroke-width: 0.2em; }"
                      )
      head.append(">")
  
      lines.append("} // digraph G")
      return "\n".join(head + lines) + "\n"
  
imports: import collections
  import contextlib
  import functools
  import html
  from typing import Any, Callable, NamedTuple, Optional
  from haiku._src import data_structures
  from haiku._src import module
  from haiku._src import utils
  import jax
  import jax.core
  from jax.experimental import pjit
  from jax.extend import linear_util
label_desc: Use image node_id 10 for calling a global function with example usage: _graph_to_dot(graph, args, outputs) and returns: unknown, out