fix more bugs

Author: majosm

pytato/analysis/__init__.py

@@ -571,13 +571,16 @@ def get_cache_key(self, expr: ArrayOrNames) -> int:  # type: ignore[override]
 
     def rec(self, expr: ArrayOrNames, *args: Any, **kwargs: Any) -> None:
         """Call the mapper method of *expr* and return the result."""
-        self.depth += 1
-        self.max_depth = max(self.max_depth, self.depth)
-
-        try:
+        if isinstance(expr, DictOfNamedArrays):
             super().rec(expr, *args, **kwargs)
-        finally:
-            self.depth -= 1
+        else:
+            self.depth += 1
+            self.max_depth = max(self.max_depth, self.depth)
+
+            try:
+                super().rec(expr, *args, **kwargs)
+            finally:
+                self.depth -= 1
 
 
 def get_max_node_depth(outputs: Union[Array, DictOfNamedArrays]) -> int:

test/test_pytato.py

@@ -769,7 +769,10 @@ def test_nodemaxdepthmapper():
     from pytato.analysis import get_max_node_depth
 
     x = pt.make_placeholder("x", shape=(10, 4), dtype=np.float64)
-    for _ in range(9):
+
+    assert get_max_node_depth(x) == 0
+
+    for _ in range(10):
         x = x + 1
 
     assert get_max_node_depth(x) == 10