Is there anyway to analyze activations in flax?

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):

```python
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.