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
google · Aug 11, 2023 · c8424b0 · c8424b0
1 parent d05aeab
commit c8424b0
Show file tree

Hide file tree

Showing 3 changed files with 557 additions and 180 deletions.
diff --git a/flax/linen/__init__.py b/flax/linen/__init__.py
@@ -92,6 +92,7 @@
     enable_named_call as enable_named_call,
     init_with_output as init_with_output,
     init as init,
+    intercept_methods as intercept_methods,
     merge_param as merge_param,
     nowrap as nowrap,
     override_named_call as override_named_call,