Lux.jl Named Layers for Container Types

Currently containers can only take tuples, and they set the named to layer_1, layer_2.... https://github.com/avik-pal/Lux.jl/pull/44 was an initial prototype for this. But we need to do it for all container layers implemented in Lux.

Jun 29 '22 03:06 avik-pal

It would be nice if one could solve this issue to get type stable containers. What was missing in that PR besides applybranching and applypairwisefusion? Maybe I will have some spare time to give it a try the next few days.

Aug 01 '22 17:08 lungd

In terms of the functionality, nothing was missing except applybranching, applyparallel, and applypairwisefusion.

I am curious as to what you mean by type stable containers. Currently all the container layers in lux should be type stable

Aug 02 '22 03:08 avik-pal

I tried to recreate the MWE I played with some weeks ago. I don't see the performance drop any longer (~20 additional allocations + exec time). Most likely, I actually benchmarked the construction of the chain. That would explain the additional allocations because I see a Body::Any for the construction with the current version vs no red line and 0 allocations with a NamedTuple.

Sorry for the false alarm. Finishing that PR to get rid of the red lines one sees using @code_warntype would be a good idea, though.

Aug 02 '22 08:08 lungd