def METHOD_NAME(self, x, state, **kwargs):
    # Apply layers
    prev_states = [None] * len(self.layers) if state is None else state
    next_states = []
    for layer, prev_state in zip(self.layers, prev_states):
        x, state = layer.METHOD_NAME(x, state=prev_state, **kwargs)
        next_states.append(state)
    x = self.norm(x)
    return x, next_states