def METHOD_NAME(self):
    # Slightly hacky way to implement this in a curried manner
    # (so that the function can be extracted from an instance)
    # Somewhat more sound may be to turn this into a
    # @staticmethod and grab subclasses using hydra.utils.get_class
    def fn(state):
        x = [
            _layer.METHOD_NAME(_state)
            for (_layer, _state) in zip(self.layers, state)
        ]
        x = [_x for _x in x if _x is not None]
        return torch.cat(x, dim=-1)
    return fn