complexity tracing: support pytrees

jcmgray · jcmgray · commit 41e8700edf91 · 2025-03-12T17:52:03.000-07:00
diff --git a/autoray/experimental/complexity_tracing.py b/autoray/experimental/complexity_tracing.py
@@ -240,11 +240,12 @@ def frequencies(it):
 
 
 def compute_cost_scalings(z, factor_map, print_missed=True):
+    from autoray.lazy import descend
 
     counts = {}
     missed = {}
 
-    for node in z.descend():
+    for node in descend(z):
         f = node.fn_name
 
         if f in COST_SCALINGS:
