Feature request: TimeDistributed Layer

[NeroBlackstone]

Motivation:
If Lux.Recurrence return_sequence arg is true, it will return a vector of 2d-matrix.

model = Chain(Recurrence(GRUCell(3=>32),return_sequence = true))

rng = Random.default_rng()
Random.seed!(rng, 0)
ps, st = Lux.setup(rng, model)

x = rand(Float32,3,4,5)
y1 = model(x,ps,st)[1]

@show size(y1)
@show size(y1[1])

# sequence_length
#size(y1) = (4,)
# output_size * batch_size
#size(y1[1]) = (32, 5)

Suppose we have a sequence labeling task, and the output of each time step is fed into a Dense layer (as a classifier) ​​for classification, looks like this:

Can the existing layers do this? What do you think?

I think it works like this:

model = Chain(Recurrence(GRUCell(3=>32),return_sequence = true),TimeDistributed(Dense(32=>2)))
rng = Random.default_rng()
Random.seed!(rng, 0)
ps, st = Lux.setup(rng, model)

x = rand(Float32,3,4,5)
y1 = model(x,ps,st)[1]

@show size(y1)
@show size(y1[1])

# sequence_length
#size(y1) = (4,)
# output_size * batch_size
#size(y1[1]) = (2, 5)

I have not studied TimeDistributed in Keras, and I don’t know if my understanding is correct.

Is it necessary to have a function as a parameter in order to merge the output of TimeDistributed? like:

...
TimeDistributed(Dense(32=>2);vcat)
...
# output
# sequence_length (after vcat) * batch_size
#size(y1) = (8, 5)

Do you think this feature is necessary? If so, I can implement it. Please tell me some ideas (such as existing layers that can be referenced) and I will complete it. Thank you!

[avik-pal]

This is a special case of #426 where the number of layers is 1.