Building a custom Layer in Lux.jl

That’s on purpose: it should return a tuple instead of an array. See (1) vs (1,).

And the function is shadowing instead of extending. You want in that script:

function Lux.initialparameters(rng::AbstractRNG, d::Polylayer)
     return (coeffs=d.init_weight(rng, d.degree+1),)
end

Lux.initialstates(::AbstractRNG, ::Polylayer) = NamedTuple()

Lux.statelength(d::Polylayer) = 0

function Lux.parameterlength(d::Polylayer)
    return d.degree + 1
end