Restore the typechecking flag when the body raises

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):

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