How to maintain a package that supports both Flux & Lux?

[avik-pal]

This question has come up several times. I should add an example for this in the manual. Basic sketch would be something like this:

using LuxCore

struct MyCustomLuxLayer <: AbstractLuxLayer
    ....
end

LuxCore.initialparameters(....) = ....
LuxCore.initialstates(....) = ....

function (l::MyCustomLuxLayer)(x, ps, st)
    ....
    return y, new_st
end

# This part is general, doesn't need to be defined for every specific layer
mutable struct MyCustomFluxLayer{L <: MyCustomLuxLayer}
    internal_layer::L
    ps
    st
end

function MyCustomFluxLayer(args...; kwargs...)
    layer = MyCustomLuxLayer(args...; kwargs...)
    return MyCustomFluxLayer(layer, LuxCore.setup(rng, layer)...)
end

function (l::MyCustomFluxLayer)(x)
    y, st = l.internal_layer(x, l.ps, l.st)
    l.st = st # Can also fmap over `st` and copy into existing buffers
    return y
end

In MyPackageFluxExt.jl

using Flux

Flux.@layer MyCustomFluxLayer (:ps,)

[MartinuzziFrancesco]

It's pretty cool that one can do this pretty seamlessly! However, there are many helper functions in Lux that are useful when building libraries for it. So I guess that if I had using Lux and and extension with Flux, that extension would have to load both the packages with some overhead right?

[avik-pal]

Yes you would have to load both in that case. 1 thing to note is that Flux does ship with AD libraries packaged, so its load time is generally higher. Also Lux deps are a subset (for the most part) of Flux, so there isn't much additional overhead

[avik-pal]

julia> @time using Lux
  0.540649 seconds (883.18 k allocations: 51.847 MiB, 0.97% gc time, 34.37% compilation time: 79% of which was recompilation)

julia> @time using Flux
  0.923958 seconds (1.45 M allocations: 88.196 MiB, 1.99% gc time, 10.92% compilation time: 32% of which was recompilation)