feat: towards a less boilerplate API for custom layers

[avik-pal]

using Lux, Random

struct LayerWithStepCounter <: AbstractLuxLayer end

LuxCore.initialstates(::AbstractRNG, ::LayerWithStepCounter) = (; counter=0)

function LuxCore.apply(::Type{<:LayerWithStepCounter}, l, x)
    l.counter += 1
    return x
end

struct MyCustomLayer2{A,B,C} <: AbstractLuxContainerLayer{(:a, :b, :c)}
    a::A
    b::B
    c::C
end

function LuxCore.initialparameters(rng::AbstractRNG, l::MyCustomLayer2)
    return (
        a=LuxCore.initialparameters(rng, l.a),
        b=LuxCore.initialparameters(rng, l.b),
        c=LuxCore.initialparameters(rng, l.c),
        d=randn(rng, Float32, 2, 3),
    )
end

function LuxCore.initialstates(rng::AbstractRNG, l::MyCustomLayer2)
    return (
        a=LuxCore.initialstates(rng, l.a),
        b=LuxCore.initialstates(rng, l.b),
        c=LuxCore.initialstates(rng, l.c),
        f=randn(rng, Float32, 2, 3),
        d=randn(rng, Float16, 2, 3),
    )
end

function LuxCore.apply(::Type{<:MyCustomLayer2}, model, x1, x2)
    y1 = model.f * model.a(x1)
    y2 = model.d * model.st_d' * model.c(model.b(x2))
    # y2 =  model.c(model.b(x2))
    return y1 .* y2
end

model = MyCustomLayer2(Dense(2 => 3), Dense(3 => 2), LayerWithStepCounter())
ps, st = Lux.setup(Random.default_rng(), model)

x1 = randn(Float32, 2, 4)
x2 = randn(Float32, 3, 4)

model((x1, x2), ps, st)

[codecov]

Codecov Report

❌ Patch coverage is 77.73109% with 53 lines in your changes missing coverage. Please review.
✅ Project coverage is 61.82%. Comparing base (4d19441) to head (a1c648a).

Files with missing lines	Patch %	Lines
src/layers/basic.jl	56.25%	21 Missing ⚠️
lib/LuxCore/src/stateful.jl	84.21%	15 Missing ⚠️
src/layers/embedding.jl	57.89%	8 Missing ⚠️
src/layers/normalize.jl	81.25%	6 Missing ⚠️
lib/LuxCore/src/apply.jl	84.61%	2 Missing ⚠️
lib/LuxCore/ext/LuxCoreFunctorsExt.jl	0.00%	1 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (4d19441) and HEAD (a1c648a). Click for more details.

HEAD has 30 uploads less than BASE

Flag	BASE (4d19441)	HEAD (a1c648a)
	51	21

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #1473       +/-   ##
===========================================
- Coverage   84.20%   61.82%   -22.39%     
===========================================
  Files         166      163        -3     
  Lines        6725     6758       +33     
===========================================
- Hits         5663     4178     -1485     
- Misses       1062     2580     +1518     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

Benchmark Results (ASV)

	main	b6fdfbc...	main / b6fdfbc...
basics/overhead	0.13 ± 0.0042 μs	0.125 ± 0.0013 μs	1.04 ± 0.036
time_to_load	0.841 ± 0.0051 s	0.842 ± 0.0027 s	0.999 ± 0.0069

Benchmark Plots

A plot of the benchmark results have been uploaded as an artifact to the workflow run for this PR.
Go to "Actions"->"Benchmark a pull request"->[the most recent run]->"Artifacts" (at the bottom).