Does FLAX have an API similar to PyTorch register_forward_hook and register_forward_hook defined in torch.nn.Module class?

[2catycm]

I think supporting hooking the __call__ function is very important, not only for debugging and visualization, but also for dynamic network structure modification.
But I cannot find the concept of hook in flax's documentation.

See FLAX's latest documentation on Module class and PyTorch's register_forward_hook.

[2catycm]

OK, now I find it, it is called interceptor https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.intercept_methods

[2catycm]

It is more general than hook. It adopts the concept of decorator in Python. Looks awesome.