flax
9d453659 - Fix bound method auto-unbinding for NNX transforms

Commit
192 days ago
Fix bound method auto-unbinding for NNX transforms - Add _resolve_bound_callable helper to detect and unbind Module methods - Update all transforms (remat, jit, grad, value_and_grad, vmap, etc.) to handle bound methods - Add comprehensive test suite for bound method functionality - Resolves issue #5053 TraceContextError with Module methods
Author
Parents
Loading