class Feeder(Py_Module):
    def __init__(self, N: int, batch_size: int, dtype):
        Py_Module.__init__(self)
        self.name = "py_Feeder"
        self.N = N
        self.batch_size = batch_size
        self.cnt = 0

        step = self.create_task("step")
        inpt = self.create_socket_in(step, "input", batch_size * N, dtype)
        out = self.create_socket_out(step, "output", N, dtype)
        self.create_codelet(step, lambda slf, lsk, fid: slf.step(lsk[inpt], lsk[out]))

    def step(self, x: ndarray, y: ndarray):
        if self.cnt >= self.batch_size:
            print("Feeder overflow")
            self.toggle_done()
            return 0
        y[:] = x[:, self.cnt * self.N: (self.cnt + 1) * self.N]
        self.cnt += 1
        if self.cnt >= self.batch_size:
            self.toggle_done()
        return 0