Adding parameters to custom layers in Flux

General Usage
1

At the output of my Flux model, I need to un-normalize one of the outputs. I only need to store a single variable value somewhere (I think the best place is in a custom layer), but I don’t know how to store an ancillary variable in the struct and have it work with the rest of the custom layer code.

Can someone kindly point me in the right direction? The code snippet mimics the Dense layer definition except for the parts for which I’ve added comments.

struct Layer{F,S<:AbstractArray,T<:AbstractArray}
    W::S
    b::T
    σ::F
    offset # desired ancillary variable used for normalization
end

Layer(W, b) = Layer(W, b, identity) # want to also set the value of Layer.offset here somehow...

function Layer(in::Integer, out::Integer, σ=identity)
    return Layer(randn(out, in), randn(out), σ)
end

Flux.@functor Layer

function (a::Layer)(x::AbstractArray)
    a.σ(a.W * x .+ a.b) .+ [ i==1 ? a.offset : 0 for i in 1:length(a.b) ] # undoes normalization in first output
end

my_layer = Layer(3, 3, softplus)

A working example of a custom layer without any ancillary values is seen here.