Count flops in array initialisations

coneoproject · Jul 6, 2018 · 8563a56 · 8563a56
1 parent 5bb1b30
commit 8563a56
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/coffee/visitors/utilities.py b/coffee/visitors/utilities.py
@@ -337,6 +337,13 @@ def visit_object(self, o, *args, **kwargs):
     def visit_list(self, o, *args, **kwargs):
         return sum(self.visit(e) for e in o)
 
+    def visit_ArrayInit(self, o, *args, **kwargs):
+        vals = o.values
+        if isinstance(vals, np.ndarray):
+            return sum(self.visit(vals[i]) for i in np.ndindex(vals.shape))
+        else:
+            return self.visit(vals)
+
     def visit_Node(self, o, *args, **kwargs):
         ops, _ = o.operands()
         return sum(self.visit(op) for op in ops)