From 1ea280c70d53dfb9ad39cbd1a3ff119fae928257 Mon Sep 17 00:00:00 2001
From: Ruslan Rakhimov <risonyo@gmail.com>
Date: Mon, 31 Aug 2026 23:10:49 +0300
Subject: [PATCH] fix: restore typechecking flag when body raises

disable_typechecking() sets the process-global jaxtyping_disable flag and
restores it after the yield. Without try/finally an exception escaping the
body skips the restore, so the flag stays True for the life of the process
and every @at.typecheck stops checking, silently.

No behaviour change on the happy path: contextlib resumes the generator at
the yield in both cases. Nesting still restores to the outer value, since
initial is captured per entry.

Adds one regression test.
---
 src/openpi/shared/array_typing.py      |  6 ++++--
 src/openpi/shared/array_typing_test.py | 13 +++++++++++++
 2 files changed, 17 insertions(+), 2 deletions(-)
 create mode 100644 src/openpi/shared/array_typing_test.py

diff --git a/src/openpi/shared/array_typing.py b/src/openpi/shared/array_typing.py
index 569eafef14..2fa1fc02cd 100644
--- a/src/openpi/shared/array_typing.py
+++ b/src/openpi/shared/array_typing.py
@@ -57,8 +57,10 @@ def typecheck(t: T) -> T:
 def disable_typechecking():
     initial = config.jaxtyping_disable
     config.update("jaxtyping_disable", True)  # noqa: FBT003
-    yield
-    config.update("jaxtyping_disable", initial)
+    try:
+        yield
+    finally:
+        config.update("jaxtyping_disable", initial)
 
 
 def check_pytree_equality(*, expected: PyTree, got: PyTree, check_shapes: bool = False, check_dtypes: bool = False):
diff --git a/src/openpi/shared/array_typing_test.py b/src/openpi/shared/array_typing_test.py
new file mode 100644
index 0000000000..7b54648468
--- /dev/null
+++ b/src/openpi/shared/array_typing_test.py
@@ -0,0 +1,13 @@
+from jaxtyping import config
+import pytest
+
+from openpi.shared import array_typing as at
+
+
+def test_disable_typechecking_restores_flag_on_exception():
+    initial = config.jaxtyping_disable
+
+    with pytest.raises(RuntimeError, match="boom"), at.disable_typechecking():
+        raise RuntimeError("boom")
+
+    assert config.jaxtyping_disable == initial