perf(typing): replace inspect.stack with bounded frame walk

src/openpi/shared/array_typing.py

@@ -1,6 +1,7 @@
 import contextlib
 import functools as ft
 import inspect
+import sys
 from typing import TypeAlias, TypeVar, cast
 
 import beartype
@@ -32,12 +33,24 @@
 
 
 def _check_dataclass_annotations(self, typechecker):
-    if not any(
-        frame.frame.f_globals.get("__name__") in {"jax._src.tree_util", "flax.nnx.transforms.compilation"}
-        for frame in inspect.stack()
-    ):
-        return _original_check_dataclass_annotations(self, typechecker)
-    return None
+    # Prefer a bounded sys._getframe walk over inspect.stack() — the latter
+    # materializes full frame info and can dominate infer latency in deep stacks.
+    # Logic preserved: skip jaxtyping checks while walking through JAX/Flax
+    # tree unflatten / compilation helpers (jaxtyping#277).
+    try:
+        frame = sys._getframe(2)
+        for _ in range(20):
+            if frame.f_globals.get("__name__") in {
+                "jax._src.tree_util",
+                "flax.nnx.transforms.compilation",
+            }:
+                return None
+            if frame.f_back is None:
+                break
+            frame = frame.f_back
+    except (ValueError, AttributeError):
+        pass
+    return _original_check_dataclass_annotations(self, typechecker)
 
 
 jaxtyping._decorator._check_dataclass_annotations = _check_dataclass_annotations  # noqa: SLF001