Merge pull request #193 from fastmachinelearning/feature/analysis-subgraph-traversal

Add apply_to_subgraph parameter to ModelWrapper.analysis()

Author: Yaman Umuroglu

src/qonnx/core/modelwrapper.py

@@ -28,6 +28,7 @@
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 import copy
+import inspect
 import onnx
 import onnx.helper as oh
 import onnx.numpy_helper as np_helper
@@ -125,9 +126,13 @@ def save(self, filename):
         """Saves the wrapper ONNX ModelProto into a file with given name."""
         onnx.save(self._model_proto, filename)
 
-    def analysis(self, analysis_fxn):
+    def analysis(self, analysis_fxn, apply_to_subgraphs=False):
         """Runs given anaylsis_fxn on this model and return resulting dict."""
-        return analysis_fxn(self)
+        if apply_to_subgraphs == True:
+            assert "apply_to_subgraphs" in inspect.signature(analysis_fxn), "analysis_fxn must have 'apply_to_subgraphs' argument when apply_to_subgraphs == True"
+            return analysis_fxn(self, apply_to_subgraphs)
+        else:
+            return analysis_fxn(self)
 
     def transform_subgraphs(self, transformation, make_deepcopy=True, cleanup=True, apply_to_subgraphs=False, use_preorder_traversal=True):
         """Applies given Transformation to all subgraphs of this ModelWrapper instance.

tests/core/test_subgraph_traversal.py

@@ -241,3 +241,17 @@ def test_traversal_nested(tree, cleanup, make_deepcopy):
 
     check_all_visted_once(tree, transform.dummy_transform)
     check_all_subgraphs_transformed(t_model.model.graph)
+
+def dummy_analysis_fxn(model_wrapper):
+    """
+    A dummy analysis function that simply returns the model wrapper.
+    This is used to test that analysis functions are called correctly.
+    """
+    d = {}
+    return d
+
+@pytest.mark.xfail(reason="Analysis functions require apply_to_subgraphs when traversing subgraphs")
+def test_analysis_fxn_without_apply_to_subgraphs_fails():
+    # Check that an analysis function fails when apply_to_subgraphs is False
+    model = make_subgraph_model(("top", [("sub1", []), ("sub2", [])]))
+    model.analysis(dummy_analysis_fxn, apply_to_subgraphs=True)