fix: relax constraints and fix nnsmith.fuzz

nnsmith/cli/fuzz.py

@@ -3,7 +3,6 @@
 from pathlib import Path
 
 import hydra
-import networkx as nx
 from omegaconf import DictConfig
 
 from nnsmith.backends.factory import BackendFactory
@@ -82,7 +81,7 @@ def make_testcase(self, seed) -> TestCase:
 
         model = self.ModelType.from_schedule(schedule)
         if self.cfg["debug"]["viz"]:
-            model.attach_viz(nx.nx_agraph.to_agraph(fixed_graph).to_string())
+            model.attach_viz(fixed_graph)
 
         model.refine_weights()  # either random generated or gradient-based.
         oracle = model.make_oracle()

nnsmith/materialize/__init__.py

@@ -15,7 +15,7 @@
 from nnsmith.abstract.op import AbsOpBase, Constant, Input
 from nnsmith.abstract.tensor import AbsTensor
 from nnsmith.error import SanityCheck
-from nnsmith.util import viz_dot
+from nnsmith.util import HAS_PYGRAPHVIZ, viz_dot
 
 
 def framework_operator_impl(
@@ -208,8 +208,9 @@ def operators() -> List[Type[AbsOpBase]]:
     def add_seed_setter() -> None:
         pass
 
-    def attach_viz(self, dotstring: str) -> None:
-        self.dotstring = dotstring
+    def attach_viz(self, graph: nx.MultiDiGraph) -> None:
+        if HAS_PYGRAPHVIZ:
+            self.dotstring = nx.nx_agraph.to_agraph(graph).to_string()
 
     def dump_viz(self, path: PathLike) -> None:
         viz_dot(self.dotstring, path)

nnsmith/util.py

@@ -1,12 +1,26 @@
 import os
 import random
 import shutil
-from typing import Callable, Dict, List, Union
+from typing import Callable, Dict, List
 
 import numpy as np
-import pygraphviz as pgv
 from termcolor import colored
 
+try:
+    import pygraphviz as pgv
+
+    HAS_PYGRAPHVIZ = True
+except ImportError:
+    import warnings
+
+    warnings.warn(
+        "Install pygraphviz for visualization: https://pygraphviz.github.io/documentation/stable/install.html\n"
+        "Currently graph visualization is not enabled."
+    )
+    pgv = None
+    HAS_PYGRAPHVIZ = False
+
+
 from nnsmith.logging import VIZ_LOG
 
 
@@ -101,7 +115,7 @@ def _check_dot_install():
     return True
 
 
-def viz_dot(dotobj: Union[pgv.AGraph, str], filename: str = None):
+def viz_dot(dotobj, filename: str = None):
     if _check_dot_install():
         if filename is None:
             filename = f"graph.png"

requirements/core.txt

@@ -1,9 +1,9 @@
 z3-solver>=4.11.0
 networkx>=2.6.3
 hydra-core>=1.1.0
-pygraphviz
 rich
 termcolor
 multipledispatch
 gitpython
 appdirs
+pygraphviz

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = nnsmith
-version = 0.1.0.a7
+version = 0.1.0.a8
 description = "Automatic DNN generation for fuzzing and more."
 long_description = file: README.md
 long_description_content_type = text/markdown
@@ -17,7 +17,6 @@ install_requires =
     z3-solver>=4.11.0
     networkx>=2.6.3
     hydra-core>=1.1.0
-    pygraphviz
     rich
     termcolor
     multipledispatch
@@ -40,4 +39,4 @@ console_scripts =
     nnsmith.model_gen = nnsmith.cli.model_gen:main
     nnsmith.model_exec = nnsmith.cli.model_exec:main
     nnsmith.dtype_test = nnsmith.cli.dtype_test:main
-    nnsmith.fuzz = nnsmith.cli.dtype_test:main
+    nnsmith.fuzz = nnsmith.cli.fuzz:main