Experimental

eqxvision.experimental.intermediate_layer_getter(model: eqx.Module, get_target_layers: Callable) -> eqx.Module

Wraps intermediate layers of a model for accessing intermediate activations. Based on a discussion here.

Info

Only supports storing the result of the most recent call. So, if the forward utilises the same layer multiple times, the returned intermediate value will be of the last call

Arguments:

• model: A PyTree representing the neural network model
• get_target_layers: A callable function which returns a sequence of layers from the model

Returns: The returned model will now return a tuple with

0. The final output of `model`
1. An ordered list of intermediate activations