Add skip_if, a step guard with non-finite and large-step conditions

[jessegrabowski]

A single NaN batch ends a run today: it reaches the parameters and every Adam moment, and nothing recovers. skip_if wraps a rule and throws away any step its condition rejects, leaving the parameters and the optimizer state as they were. Exhausting the tolerance raises rather than accepting the poisoned update the way optax does; a run that has skipped that many steps in a row has diverged, not hit a bad batch.

Two conditions ship: nonfinite, aliased apply_if_finite, and large_step, which thresholds on step norm and so fires while the numbers are still finite. large_step isn't redundant with clip_by_global_norm. Clipping rescales by max_norm / norm, which is zero once the norm is inf, so the poisoned coordinate becomes inf * 0 and every healthy one is multiplied away; one inf destroys its own parameter and cancels the step for all the others. The guard covers only what the rule writes, so batch-norm running statistics, folded in by compile_train outside it, still get poisoned.

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 97.52%. Comparing base (c1afd79) to head (4d6a905).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #107      +/-   ##
==========================================
+ Coverage   97.47%   97.52%   +0.04%     
==========================================
  Files          60       61       +1     
  Lines        3012     3069      +57     
==========================================
+ Hits         2936     2993      +57     
  Misses         76       76              

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.