code: def abstract_to_dot(fun: Callable[..., Any]) -> Callable[..., str]:
      """Converts a function using Haiku modules to a dot graph.
  
      Same as :func:`to_dot` but uses JAX's abstract interpretation
      machinery to evaluate the function without requiring concrete inputs.
      Valid inputs for the wrapped function include
      :class:`jax.ShapeDtypeStruct`.
  
      :func:`abstract_to_dot` does not support data-dependent
      control-flow, because no concrete values are provided to the function.
  
      Args:
        fun: A function using Haiku modules.
  
      Returns:
        A function that returns the source code string to a graphviz graph
        describing the operations executed by the given function clustered by Haiku
        module.
  
      See Also:
        :func:`to_dot`: Generates a graphviz graph using concrete inputs.
      """
  
      @functools.wraps(fun)
      def wrapped_fun(*args) -> str:
          dot_out = ""
  
          # eval_shape cannot evaluate functions which return str, as str is not a
          # valid JAX types.
          # The following function extracts the created dot string during the
          # abstract evaluation.
          def dot_extractor_fn(*inner_args):
              nonlocal dot_out
              dot_out = to_dot(fun)(*inner_args)
  
          jax.eval_shape(dot_extractor_fn, *args)
          assert (
              dot_out
          ), "Failed to extract dot graph from abstract evaluation"
          return dot_out
  
      return wrapped_fun
  
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 2 for calling a global function with example usage: abstract_to_dot(fun) and returns: wrapped_fun, dot_out