Module Interceptor.
zpolina opened this issue · comments
This FLIP discusses adding a new module interceptor API into calls to Flax modules to add functionality for various purposes. For example, interceptors would make it possible to change the input of a Module call or to perturb its output. Interceptors can also be used to access or modify the gradients that pass through a Module and expose those.
@jheek ended up implementing custom solution that did not involve hooks, and shared it offline. I am therefore closing this issue.
@marcvanzee this is not super helpful. How does this custom solution work?
I still believe a module interceptor would be very valuable.