move count_ops to common tvm.relay.testing

python/tvm/relay/testing/__init__.py

@@ -17,6 +17,7 @@
 #pylint: disable=invalid-name
 """Utilities for testing and benchmarks"""
 from __future__ import absolute_import as _abs
+import collections
 import numpy as np
 
 import tvm
@@ -135,3 +136,18 @@ def check_grad(func, inputs=None, eps=1e-6, atol=1e-5, rtol=1e-3, scale=None, me
 
 def rand(dtype, *shape):
     return tvm.nd.array(np.random.rand(*shape).astype(dtype))
+
+
+def count_ops(expr):
+    """count number of times a given op is called in the graph"""
+    class OpCounter(tvm.relay.ExprVisitor):
+        def visit_call(self, call):
+            if hasattr(call, 'op'):
+                self.node_counter[call.op.name] += 1
+            return super().visit_call(call)
+        def count(self, expr):
+            self.node_set = {}
+            self.node_counter = collections.Counter()
+            self.visit(expr)
+            return self.node_counter
+    return OpCounter().count(expr)

tests/python/relay/test_pass_gradient.py

@@ -25,24 +25,10 @@
 from tvm.relay import create_executor, transform
 from tvm.relay.transform import gradient
 from tvm.relay.prelude import Prelude
-from tvm.relay.testing import add_nat_definitions, make_nat_expr, run_infer_type, check_grad, rand
+from tvm.relay.testing import add_nat_definitions, make_nat_expr, run_infer_type, check_grad, rand, count_ops
 import tvm.relay.op as op
 
 
-def count_ops(expr):
-    class OpCounter(tvm.relay.ExprVisitor):
-        def visit_call(self, expr):
-            if hasattr(expr, 'op'):
-                self.node_counter[expr.op.name] += 1
-            return super().visit_call(expr)
-        def count(self, expr):
-            self.node_set = {}
-            self.node_counter = collections.Counter()
-            self.visit(expr)
-            return self.node_counter
-    return OpCounter().count(expr)
-
-
 def test_id():
     shape = (10, 10)
     dtype = 'float32'