Restore the typechecking flag when the body raises

[rakhimovv]

disable_typechecking() sets the process-global jaxtyping_disable flag and restores it after
the yield. With no 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 in the codebase stops checking —
nothing raised, nothing logged.

initial = config.jaxtyping_disable
config.update("jaxtyping_disable", True)
yield                                        # if this raises, the next line never runs
config.update("jaxtyping_disable", initial)

Where this actually bites. Not the training scripts — training/checkpoints.py has no
try/except and scripts/train.py calls restore_state at top level, so a failed restore ends
the process and the flag never outlives it. It bites in any process that catches and continues:

• pytest — a raise inside the context manager disables typechecking for the rest of the session,
  so later tests silently stop being typechecked
• a REPL or notebook, which is the same environment class as array_typing._check_dataclass_annotations breaks inside REPLs #596 / Fix array_typing breakage in REPL #604

No behaviour change on the happy path. contextlib._GeneratorContextManager.__exit__ resumes the
generator at the yield in both cases; return/break out of the body already route through the
normal-exit path. Nesting still restores to the outer value, since initial is captured per entry.
This does not make the flag thread-safe — it is still process-global and restore is last-writer-wins;
this only fixes the leak on exception.

Test. One regression test, in src/openpi/shared/array_typing_test.py alongside the existing
download_test.py / image_tools_test.py / normalize_test.py. It fails on main and passes with
the fix. It compares against the flag's initial value rather than asserting a literal, so it is
unaffected by JAXTYPING_DISABLE in the environment. Happy to drop it if you would rather keep the
change to the three lines.

CONTRIBUTING.md asks for an issue or discussion first — glad to move this to one if you prefer; it
seemed small enough to send directly.