Restore the global RKG key when a transform raises

[cemde]

Bug

Inside a transform, _wrap_fn swaps the global RKG.key for the traced key so that layers constructed within the transform draw from the traced stream, then swaps it back:

_old_key, RKG.key = RKG.key, kwargs["__RKG"].key
_r = fn(*args, **kwargs, _is_root=False)
RKG.key = _old_key          # plain statement, not a finally

If fn raises, the restore never runs and the tracer stays in module-level state.

Impact

pcx.RKG is the default argument of every layer and Vode constructor, so a leaked tracer breaks every subsequent random draw in the process with UnexpectedTracerError, including in code that never touched a transform. One failed pxf.jit call is enough. Under pytest, one failing test breaks every later test that constructs a layer.

Fix

try/finally around the call, at both swap sites:

try:
    _r = fn(*args, **kwargs, _is_root=False)
finally:
    RKG.key = _old_key

The success path is unchanged; return _r, kwargs stays outside the try. The flow transforms (scan, while_loop, cond, switch) inherit _BaseTransform and route through the same closures, so both guarding tests are covered by these two edits.

Nesting restores correctly in LIFO order: for jit(value_and_grad(f)) the inner finally restores the outer transform's tracer, and the outer finally restores the concrete key.

Follow-up, not fixed here

Vmap._t has the same shape of bug at a different site: it splits the global key over the batch axis and merges it back with the merge outside any try, so a raising pxf.vmap leaves the key with an extra leading axis. Concrete rather than traced, so this PR's guarding test (which asserts "not a tracer") passes while the process is still poisoned. Different function, different fix, filed as #91.

Closes #71

[liukidar]

lgtm