Broadcast Layer

[avik-pal]

A simple implementation:

struct BroadcastLayer{T <: NamedTuple} <: AbstractExplicitContainerLayer{(:layers,)}
    layers::T
end

function BroadcastLayer(layers...)
    for l in layers
        if !iszero(statelength(l))
            throw(ArgumentError("Stateful layer `$l` are not supported for `BroadcastLayer`."))
        end
    end
    names = ntuple(i -> Symbol("layer_$i"), length(layers))
    return BroadcastLayer(NamedTuple{names}(layers))
end

BroadcastLayer(; kwargs...) = BroadcastLayer(connection, (; kwargs...))

function (m::BroadcastLayer)(x, ps, st::NamedTuple{names}) where {names}
    results = (first ∘ Lux.apply).(values(m.layers), x, values(ps), values(st))
    return results, st
end

Base.keys(m::BroadcastLayer) = Base.keys(getfield(m, :layers))

Originally posted by @avik-pal in #282 (comment)

[NeroBlackstone]

BroadcastLayer(; kwargs...) = BroadcastLayer(connection, (; kwargs...))

What is the role of kwargs and connection in this line?
Do I need to add the implementation of BroadcastLayer(connection, (; kwargs...))?

[NeroBlackstone]

I still can't understand what BroadcastLayer does when there are multiple layers passed in.

1. If there are n layers in the BroadcastLayer, and length of x is 1, x will be copied n times and passed to every internal layers respectively.
2. If there are n layers in the BroadcastLayer, and length of x is n, x_i will be passed to the corresponding layer_i.
3. If there is only 1 layer in the BroadcastLayer, and length of x is n, x_i will be passed to layer respectively.
4. If there are n layers in the BroadcastLayer, and length of x is m, m != n, m>1 and n>1, throw a error.

Is my understanding correct and is it necessary to document this?

[NeroBlackstone]

@avik-pal Hi, Could you help me confirm that my understanding is correct? If its correct I will open a PR soon

[avik-pal]

Honesty with the introduction of @compact, it might make sense not to have them inbuilt. Broadcasting is effectively:

@compact(; layer = <your layer>) do x::Union{NTuple{<:AbstractArray}, AbstractVector{<:AbstractArray}
    @return map(layer, x)
end

[NeroBlackstone]

Honesty with the introduction of @compact, it might make sense not to have them inbuilt. Broadcasting is effectively:

@compact(; layer = <your layer>) do x::Union{NTuple{<:AbstractArray}, AbstractVector{<:AbstractArray}
    @return map(layer, x)
end

@avik-pal Thank you very much! This is useful!
But I still think inbuilt a broadcast layer is useful, this can reduce some of template code.
Or we need to minimally document this code in RNN related documentation.
What do you think, implement a broadcast layer or just simply document this example? I can open a PR.