How to List Full Model Topology, including Activation Layers

[oelghaza]

I am looking for a way to reconstruct my NN model while using the new NNX API. If I have a basic MLP with Activations between Linear layers, I am not able to retrieve any information on the Activation layers.

I understand this is due to them being JAX functions rather than NNX Modules. Is there a way for me to perform model surgery and retrieve information about the activation layers of a saved model — without explicitly creating a Flax Module for each activation layer?

Ie if my model has the following architecture:

-Linear1
-Activation1
-Linear2
-Activation2

Using the model surgery techniques in the guides(display(), iter_children(), iter_modules(), saving/restoring checkpoint) I can only retrieve the following:

-Linear1
-Linear2

[cgarciae]

Hey! Flax NNX will print all instance attributes of a Module so if you refactor your code such that the activation functions are attributes you should be able to find them with nnx.iter_graph.