flax
c8424b04 - Add a Haiku-like methods interceptor API.

Commit
2 years ago
Add a Haiku-like methods interceptor API. Note: 1. Methods decorated by nn.nowrap could not be intercepted. 2. Module dataclass fields and dunder methods, including __eq__, __repr__, __init__, __hash__, and __post__init__ could not be intercepted. 3. Module descriptors could not be intercepted. 4. InterceptorContext currently contains the module object, method name, and original method. More fields could be added in the future. PiperOrigin-RevId: 556064468
Author
Committer
Parents
Loading