Apply a stepwise model over a sequence, maintaining state. For RNNs

def Recurrent(step_model):
    """Apply a stepwise model over a sequence, maintaining state. For RNNs"""
    ops = step_model.ops

    def recurrent_fwd(seqs, drop=0.0):
        lengths = [len(X) for X in seqs]
        X, size_at_t, unpad = ops.square_sequences(seqs)
        Y = ops.allocate((X.shape[0], X.shape[1], step_model.nO))
        cell_drop = ops.get_dropout_mask((len(seqs), step_model.nO), 0.0)
        hidden_drop = ops.get_dropout_mask((len(seqs), step_model.nO), 0.0)
        out_drop = ops.get_dropout_mask((len(seqs), step_model.nO), 0.0)
        backprops = [None] * max(lengths)
        state = step_model.weights.get_initial_state(len(seqs))
        for t in range(max(lengths)):
            state = list(state)
            size = size_at_t[t]
            Xt = X[t, :size]
            state[0] = state[0][:size]
            state[1] = state[1][:size]
            if cell_drop is not None:
                state[0] *= cell_drop
            if hidden_drop is not None:
                state[1] *= hidden_drop
            inputs = (state, Xt)
            (state, Y[t, :size]), backprops[t] = step_model.begin_update(inputs)
            if out_drop is not None:
                Y[t, :size] *= out_drop
        outputs = unpad(Y)

        def recurrent_bwd(d_outputs, sgd=None):
            dY, size_at_t, unpad = step_model.ops.square_sequences(d_outputs)
            d_state = [
                step_model.ops.allocate((dY.shape[1], step_model.nO)),
                step_model.ops.allocate((dY.shape[1], step_model.nO)),
            ]
            updates = {}

            def gather_updates(weights, gradient, key=None):
                updates[key] = (weights, gradient)

            dX = step_model.ops.allocate(
                (dY.shape[0], dY.shape[1], step_model.weights.nI)
            )
            for t in range(max(lengths) - 1, -1, -1):
                if out_drop is not None:
                    dY[t] *= out_drop
                d_state_t, dXt = backprops[t]((d_state, dY[t]), sgd=gather_updates)
                d_state[0][: d_state_t[0].shape[0]] = d_state_t[0]
                d_state[1][: d_state_t[1].shape[0]] = d_state_t[1]
                dX[t, : dXt.shape[0]] = dXt
                if cell_drop is not None:
                    d_state[0] *= cell_drop
                if hidden_drop is not None:
                    d_state[1] *= hidden_drop
            d_cell, d_hidden = d_state
            step_model.weights.d_initial_cells += d_cell.sum(axis=0)
            step_model.weights.d_initial_hiddens += d_hidden.sum(axis=0)
            if sgd is not None:
                for key, (weights, gradient) in updates.items():
                    sgd(weights, gradient, key=key)
            return unpad(dX)

        return outputs, recurrent_bwd

    model = wrap(recurrent_fwd, step_model)
    model.nO = step_model.nO
    return model