Don't revisit shared subjaxprs in jaxpr_util.pprof_equation_profile.

It is probably a more useful default behavior not to implicitly inline everything.

PiperOrigin-RevId: 769452443

Author: hawkinsp

jax/_src/jaxpr_util.py

@@ -33,11 +33,23 @@
 zip, unsafe_zip = util.safe_zip, zip
 
 
-def all_eqns(jaxpr: core.Jaxpr) -> Iterator[tuple[core.Jaxpr, core.JaxprEqn]]:
+def _all_eqns(
+    jaxpr: core.Jaxpr, visited: set[core.Jaxpr] | None,
+) -> Iterator[tuple[core.Jaxpr, core.JaxprEqn]]:
   for eqn in jaxpr.eqns:
     yield (jaxpr, eqn)
   for subjaxpr in core.subjaxprs(jaxpr):
-    yield from all_eqns(subjaxpr)
+    if visited is None:
+      yield from _all_eqns(subjaxpr, visited)
+    elif subjaxpr not in visited:
+      visited.add(subjaxpr)
+      yield from _all_eqns(subjaxpr, visited)
+
+def all_eqns(
+    jaxpr: core.Jaxpr, revisit_inner_jaxprs: bool = True
+) -> Iterator[tuple[core.Jaxpr, core.JaxprEqn]]:
+  yield from _all_eqns(jaxpr, None if revisit_inner_jaxprs else set())
+
 
 def collect_eqns(jaxpr: core.Jaxpr, key: Callable):
   d = defaultdict(list)
@@ -206,7 +218,7 @@ def pprof_equation_profile(jaxpr: core.Jaxpr) -> bytes:
   """
   d = Counter(
       (eqn.source_info.traceback, eqn.primitive)
-      for _, eqn in all_eqns(jaxpr)
+      for _, eqn in all_eqns(jaxpr, revisit_inner_jaxprs=False)
   )
   return _pprof_profile(d)