May I ask how to obtain the output of each layer and the activation values when i use JAX?

[sglucas]

May I ask how to obtain the output of each layer and the activation values when i use JAX? I hope you could provide a simple example.

Best,
Lucas

[jakevdp]

Thanks for the question! In general the answer would be to write your function in a way such that the intermediate values are returned. If you're using a JAX-based neural network framework like flax, haiku, equinox, etc. then there is probably a more specific way to do what you ask, but I'd suggest asking in the appropriate forum for the specific tool.

[sglucas]

Hi @jakevdp Thanks for your answer. I'm using flax to build the network.

[jakevdp]

Sounds good - the folks here would know the answer: https://github.com/google/flax/discussions