Add node type counter

pytato/analysis/__init__.py

@@ -26,7 +26,7 @@
 """
 
 from typing import (Mapping, Dict, Union, Set, Tuple, Any, FrozenSet,
-                    TYPE_CHECKING)
+                    Type, TYPE_CHECKING)
 from pytato.array import (Array, IndexLambda, Stack, Concatenate, Einsum,
                           DictOfNamedArrays, NamedArray,
                           IndexBase, IndexRemappingBase, InputArgumentBase,
@@ -413,6 +413,49 @@ def get_num_nodes(outputs: Union[Array, DictOfNamedArrays]) -> int:
 # }}}
 
 
+# {{{ NodeTypeCountMapper
+
+@optimize_mapper(drop_args=True, drop_kwargs=True, inline_get_cache_key=True)
+class NodeTypeCountMapper(CachedWalkMapper):
+    """
+    Counts the number of nodes of a given type in a DAG.
+
+    .. attribute:: counts
+
+       Dictionary mapping node types to number of nodes of that type.
+    """
+
+    def __init__(self) -> None:
+        from collections import defaultdict
+        super().__init__()
+        self.counts = defaultdict(int)
+
+    def get_cache_key(self, expr: ArrayOrNames) -> int:
+        return id(expr)
+
+    def post_visit(self, expr: Any) -> None:
+        if type(expr) not in self.counts:
+            self.counts[type(expr)] = 0
+        self.counts[type(expr)] += 1
+
+
+def get_num_node_types(outputs: Union[Array, DictOfNamedArrays]) -> Dict[Type, int]:
+    """
+    Returns a dictionary mapping node types to node count for that type
+    in DAG *outputs*.
+    """
+
+    from pytato.codegen import normalize_outputs
+    outputs = normalize_outputs(outputs)
+
+    ncm = NodeTypeCountMapper()
+    ncm(outputs)
+
+    return ncm.counts
+
+# }}}
+
+
 # {{{ CallSiteCountMapper
 
 @optimize_mapper(drop_args=True, drop_kwargs=True, inline_get_cache_key=True)