Is there anyway to analyze activations in flax?

[mohamad-amin]

A common thing in deep learning research/engineering is to analyze the intermediate activations of a model. In PyTorch, this is fairly simple to do (though I think it can be even simpler):

import torch, torch.nn as nn, torch.nn.functional as F

class SimpleMLP(nn.Module):

    def __init__(self): 
        super().__init__()
        self.fc1 = nn.Linear(5, 10);
        self.fc2 = nn.Linear(10, 2)
        self.fc3 = nn.Linear(2, 1)

    def forward(self, x): 
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

model, X = SimpleMLP(), torch.randn(3, 5)
h1 = F.relu(model.fc1(X))
h2 = F.relu(model.fc2(h1))
h3 = model.fc3(h2)

How should one implement this in flax? It's possible to write multiple apply functions, but that's really not something someone debugging a model wants to do.

[cgarciae]

Pass capture_intermediate=True to apply.