testing multilayer mgu

[MartinuzziFrancesco]

No description provided.

[MartinuzziFrancesco]

maybe something along these lines

struct TestMGU{A, B}
    cells::A
    dropout_layer::B
end

Flux.@layer :expand TestMGU

function TestMGU((in_size, out_size)::Pair;
    n_layers::Int=1,
    dropout::Float64=0.0,
    kwargs...)
    cells = []

    for i in 1:n_layers
        tin_size = i == 1 ? in_size : out_size
        push!(cells, MGUCell(tin_size => out_size; kwargs...))
    end

    if dropout > 0.0
        dropout_layer = Dropout(dropout)
    else
        dropout_layer = nothing
    end

    return TestMGU(cells, dropout_layer)
end

function (mgu::TestMGU)(inp, state)
    @assert ndims(inp) == 2 || ndims(inp) == 3
    output = []
    num_layers = length(mgu.cells)

    for inp_t in eachslice(inp, dims=2)
        new_states = []
        for (idx_cell, cell) in enumerate(mgu.cells)
            new_state = cell(inp_t, state[idx_cell])
            new_states = vcat(new_states, [new_state])
            inp_t = new_state

            if mgu.dropout_layer isa Dropout && idx_cell < num_layers - 1
                inp_t = mgu.dropout_layer(inp_t)
            end
        end
        state = stack(new_states, dims=2)

        output = vcat(output, [inp_t])
    end
    output = stack(output, dims=2)
    return output, state
end