Passing on State to BatchNorm when it is integrated with nn.Squential

[paganpasta]

Hi,

This in continuation to the effort of bringing eqxvision up to date with changes in equinox (paganpasta/eqxvision#72).

One of the changes required now are to pass the state explicitly to the Normlayer's call.
Quite often the models are recursive collections of bunch of layers with BatchNorm integrated in them somewhere. For example,
ResNet utilises mainly 4 blocks where each block is implemented as a Sequential and within each sequential block there is a BatchNorm.

One solution I see is to either modify the equinox.nn.Sequential to accomodate an optional State input or simply implement a modified Sequential in the Eqxvision repository.

@patrick-kidger, I wanted to seek your advice on the best way to go about and supporting such scenarios before making any drastic changes.

[patrick-kidger]

Heyhey! I think we could probably do something like:

class StatefulLayer(eqx.Module):
    @abc.abstractmethod
    def __call__(self, x, state):
        ...

class BatchNorm(StatefulLayer):
    ...

class Sequential(eqx.Module):
    ...

    def __call__(self, x, state=sentinel):
        for layer in self.layers:
            if isinstance(layer, StatefulLayer) and state is not sentinel:
                x, state = layer(x, state)
            else:
                x = layer(x)
        if state is sentinel:
            return x
        else:
            return x, state

In other words: (a) use the same trick that LayerNorm uses to optionally pass through stateful input, and (b) add an explicit way of checking whether a layer should take an extra state input.

I'd be happy to take a PR on this to Equinox if you want. (I'd also be happy to see this appear in Eqxvision; I don't have strong feelings either way.)

[hlzl]

I think this can be closed due to #505

[patrick-kidger]

Yup! Problem solved. :)