states_type_remains_constant(::AbstractLuxLayer) = Val(true)

# Currently, this is not defined. With this dispatch, if users
# have this dispatch (current official API), it just uses that
# dispatch. So no breaking change whatsoever
function (m::AbstractLuxLayer)(x, ps, st)
    smodel = Lux.StatefulLuxLayer{states_type_remains_constant(m)}(m, ps, st)
    # This is not done via dispatch to avoid ambiguities
    xs = x isa Tuple ? x : (x,)
    res = LuxCore.apply(typeof(m), smodel, xs...)
    if constatnt_state_type
        return res, smodel.st
    else
        return res, smodel.st_any
    end
end

function (m::MyCustomLayer)(x, ps, st)
    y, st2 = m.layer_1(x, ps.layer_1, st.layer_1)
    return y, (; layer_1 = st2)
end

# Note that `model` is not MyCustomLayer. Instead, it is `StatefulLuxLayer`
LuxCore.apply(::Type{<:MyCustomLayer}, model, x) = model.layer_1(x)

# For multi input layers, we can now have
LuxCore.apply(::Type{<:MyCustomLayer}, model, x1, x2) = model.layer_1(x1)