feat: rework more layers

Author: avik-pal

src/extended_ops.jl

@@ -266,7 +266,7 @@ for (op, field) in (
     :track_stats => :track_stats,
     :train_state => :train_state,
 )
-    @eval function $(Symbol(:has_, op))(l::AbstractLuxLayer)
+    @eval function $(Symbol(:has_, op))(l)
         res = known(safe_getproperty(l, Val($(Meta.quot(field)))))
         return ifelse(res === nothing, false, res)
     end

src/layers/embedding.jl

@@ -57,32 +57,29 @@ end
 
 outputsize(e::Embedding, _, ::AbstractRNG) = (e.out_dims,)
 
-function (e::Embedding)(x::Union{Number,AbstractVector}, ps, st::NamedTuple)
+function apply(::Type{<:Embedding}, e, x::Union{Number,AbstractVector})
     @argcheck Utils.eltype(x) <: Integer
-    return ps.weight[:, x], st
+    return ps.weight[:, x]
 end
-function (e::Embedding)(x::AbstractArray, ps, st::NamedTuple)
+function apply(T::Type{<:Embedding}, e, x::AbstractArray)
     @argcheck Utils.eltype(x) <: Integer
-    y, stₙ = e(Utils.vec(x), ps, st)
-    return reshape(y, :, size(x)...), stₙ
+    y = apply(T, e, Utils.vec(x))
+    return reshape(y, :, size(x)...)
 end
-function (e::Embedding)(x::NTuple{N,T}, ps, st::NamedTuple) where {N,T}
+function apply(::Type{<:Embedding}, e, x::T...) where {T}
     @argcheck Utils.eltype(T) <: Integer
-    return ps.weight[:, x...], st
+    return ps.weight[:, x...]
 end
-function (e::Embedding)(x::NTuple{N,<:AbstractVector{T}}, ps, st::NamedTuple) where {N,T}
+
+function apply(::Type{<:Embedding}, e, x::AbstractVector{T}...) where {T}
     @argcheck Utils.eltype(T) <: Integer
     @argcheck allequal(size, x) DimensionMismatch("Input vectors must have the same shape")
-    return NNlib.gather(ps.weight, x...), st
+    return NNlib.gather(ps.weight, x...)
 end
-function (e::Embedding)(x::NTuple{N,<:AbstractArray{T}}, ps, st::NamedTuple) where {N,T}
-    @argcheck Utils.eltype(T) <: Integer
+function apply(T::Type{<:Embedding}, e, x::AbstractArray...)
     @argcheck allequal(size, x) DimensionMismatch("Input arrays must have the same shape")
-    y, stₙ = e(vec.(x), ps, st)
-    return reshape(y, :, size(first(x))...), stₙ
-end
-function (e::Embedding)(::Tuple{}, _, ::NamedTuple)
-    throw(ArgumentError("Input tuple must contain at least one element"))
+    y = apply(T, e, vec.(x)...)
+    return reshape(y, :, size(first(x))...)
 end
 
 @doc doc"""
@@ -145,10 +142,11 @@ function initialstates(::AbstractRNG, spe::SinusoidalPositionalEmbedding{T}) whe
     return (; sigmas)
 end
 
-function (spe::SinusoidalPositionalEmbedding)(x::AbstractArray, ps, st::NamedTuple)
-    y = reshape(match_eltype(spe, ps, st, x), 1, size(x)...) .* st.sigmas
-    z = vcat(sin.(y), cos.(y)) .* spe.scale
-    return z, st
+function apply(::Type{<:SinusoidalPositionalEmbedding}, spe, x::AbstractArray)
+    # XXX: restore `match_eltype` support
+    # x′ = match_eltype(spe, x)
+    y = reshape(x, 1, size(x)...) .* spe.sigmas
+    return vcat(sin.(y), cos.(y)) .* spe.scale
 end
 
 """
@@ -220,16 +218,14 @@ function initialstates(::AbstractRNG, rope::RotaryPositionalEmbedding)
     )
 end
 
-function (rope::RotaryPositionalEmbedding)(
-    x::AbstractArray{T,4}, ps, st::NamedTuple
-) where {T}
-    y = apply_rotary_embedding(x, st.cos_cache, st.sin_cache; seq_dim=3)
-    return y, st
+function apply(::Type{<:RotaryPositionalEmbedding}, rope, x::AbstractArray{T,4}) where {T}
+    return apply_rotary_embedding(x, rope.cos_cache, rope.sin_cache; seq_dim=3)
 end
 
-function (rope::RotaryPositionalEmbedding)((x, input_pos)::Tuple, ps, st::NamedTuple)
-    y = apply_rotary_embedding(x, input_pos, st.cos_cache, st.sin_cache; seq_dim=3)
-    return y, st
+function apply(
+    ::Type{<:RotaryPositionalEmbedding}, rope, x::AbstractArray{T,4}, input_pos
+) where {T}
+    return apply_rotary_embedding(x, input_pos, rope.cos_cache, rope.sin_cache; seq_dim=3)
 end
 
 ## Functional variants since Qwen3 like models tend to share the same rotary embedding

src/layers/extension.jl

@@ -76,14 +76,11 @@ function Base.show(io::IO, ::MIME"text/plain", s::SimpleChainsLayer)
     return PrettyPrinting.print_wrapper_model(io, "SimpleChainsLayer", s.lux_layer)
 end
 
-function (sc::SimpleChainsLayer)(x, ps, st)
-    y = match_eltype(sc, ps, st, x)
-    return (
-        to_array(
-            sc.to_array,
-            apply_simple_chain(sc.layer, y, ps.params, MLDataDevices.get_device(x)),
-        ),
-        st,
+function apply(::Type{<:SimpleChainsLayer}, sc, x)
+    # XXX: restore `match_eltype` support
+    # x′ = match_eltype(sc, x)
+    return to_array(
+        sc.to_array, apply_simple_chain(sc.layer, x, sc.params, MLDataDevices.get_device(x))
     )
 end
 

src/layers/normalize.jl

@@ -134,41 +134,40 @@ end
 parameterlength(l::BatchNorm) = ifelse(has_affine(l), l.chs * 2, 0)
 statelength(l::BatchNorm) = ifelse(has_track_stats(l), l.chs * 2, 0) + 1
 
-function (BN::BatchNorm)(x::AbstractArray, ps, st::NamedTuple)
+function apply(::Type{<:BatchNorm}, bn, x::AbstractArray)
     CRC.ignore_derivatives() do
-        if st.training isa Val{true}
+        if bn.training isa Val{true}
             @argcheck size(x, ndims(x)) != 1 "Batch size for BatchNorm cannot be 1 during training"
         end
     end
 
-    x′ = match_eltype(BN, ps, st, x)
-    σ = NNlib.fast_act(BN.activation, x′)
+    # XXX: restore `match_eltype` support
+    # x′ = match_eltype(bn, ps, st, x)
+    x′ = x
+    σ = NNlib.fast_act(bn.activation, x′)
     y, stats = batchnorm(
         x′,
-        safe_getproperty(ps, Val(:scale)),
-        safe_getproperty(ps, Val(:bias)),
-        safe_getproperty(st, Val(:running_mean)),
-        safe_getproperty(st, Val(:running_var)),
-        st.training,
+        safe_getproperty(bn, Val(:scale)),
+        safe_getproperty(bn, Val(:bias)),
+        safe_getproperty(bn, Val(:running_mean)),
+        safe_getproperty(bn, Val(:running_var)),
+        bn.training,
         σ,
-        convert(unwrapped_eltype(x′), BN.momentum),
-        convert(unwrapped_eltype(x′), BN.epsilon),
+        convert(unwrapped_eltype(x′), bn.momentum),
+        convert(unwrapped_eltype(x′), bn.epsilon),
     )
-    return y, update_batchnorm_state(BN, st, stats)
+    update_batchnorm_state!(bn, stats)
+    return y
 end
 
-function update_batchnorm_state(BN::BatchNorm, st::NamedTuple, stats)
-    has_track_stats(BN) && return merge(
-        st,
-        (;
-            running_mean=Utils.vec(stats.running_mean),
-            running_var=Utils.vec(stats.running_var),
-        ),
-    )
-    return st
+function update_batchnorm_state!(bn, stats)
+    has_track_stats(bn) || return nothing
+    bn.running_mean = Utils.vec(stats.running_mean)
+    bn.running_var = Utils.vec(stats.running_var)
+    return nothing
 end
 
-CRC.@non_differentiable update_batchnorm_state(::Any...)
+CRC.@non_differentiable update_batchnorm_state!(::Any...)
 
 function Base.show(io::IO, l::BatchNorm)
     print(io, "BatchNorm($(l.chs)")
@@ -280,18 +279,20 @@ end
 
 parameterlength(l::GroupNorm) = has_affine(l) ? (l.chs * 2) : 0
 
-function (GN::GroupNorm)(x::AbstractArray, ps, st::NamedTuple)
-    x′ = match_eltype(GN, ps, st, x)
-    σ = NNlib.fast_act(GN.activation, x′)
+function apply(::Type{<:GroupNorm}, gn, x::AbstractArray)
+    # XXX: restore `match_eltype` support
+    # x′ = match_eltype(GN, ps, st, x)
+    x′ = x
+    σ = NNlib.fast_act(gn.activation, x′)
     y = groupnorm(
         x′,
-        safe_getproperty(ps, Val(:scale)),
-        safe_getproperty(ps, Val(:bias)),
-        GN.groups,
+        safe_getproperty(gn, Val(:scale)),
+        safe_getproperty(gn, Val(:bias)),
+        gn.groups,
         σ,
-        convert(unwrapped_eltype(x′), GN.epsilon),
+        convert(unwrapped_eltype(x′), gn.epsilon),
     )
-    return y, st
+    return y
 end
 
 function Base.show(io::IO, l::GroupNorm)
@@ -434,35 +435,34 @@ end
 parameterlength(l::InstanceNorm) = ifelse(has_affine(l), l.chs * 2, 0)
 statelength(l::InstanceNorm) = ifelse(has_track_stats(l), l.chs * 2, 0) + 1
 
-function (IN::InstanceNorm)(x::AbstractArray, ps, st::NamedTuple)
-    x′ = match_eltype(IN, ps, st, x)
-    σ = NNlib.fast_act(IN.activation, x′)
+function apply(::Type{<:InstanceNorm}, in, x::AbstractArray)
+    # XXX: restore `match_eltype` support
+    # x′ = match_eltype(IN, ps, st, x)
+    x′ = x
+    σ = NNlib.fast_act(in.activation, x′)
     y, stats = instancenorm(
         x′,
-        safe_getproperty(ps, Val(:scale)),
-        safe_getproperty(ps, Val(:bias)),
-        safe_getproperty(st, Val(:running_mean)),
-        safe_getproperty(st, Val(:running_var)),
-        st.training,
+        safe_getproperty(in, Val(:scale)),
+        safe_getproperty(in, Val(:bias)),
+        safe_getproperty(in, Val(:running_mean)),
+        safe_getproperty(in, Val(:running_var)),
+        in.training,
         σ,
-        convert(unwrapped_eltype(x′), IN.momentum),
-        convert(unwrapped_eltype(x′), IN.epsilon),
+        convert(unwrapped_eltype(x′), in.momentum),
+        convert(unwrapped_eltype(x′), in.epsilon),
     )
-    return y, update_instancenorm_state(IN, st, stats)
+    update_instancenorm_state!(in, stats)
+    return y
 end
 
-function update_instancenorm_state(IN::InstanceNorm, st::NamedTuple, stats)
-    has_track_stats(IN) && return merge(
-        st,
-        (;
-            running_mean=Utils.vec(stats.running_mean),
-            running_var=Utils.vec(stats.running_var),
-        ),
-    )
-    return st
+function update_instancenorm_state!(in::InstanceNorm, stats)
+    has_track_stats(in) || return nothing
+    in.running_mean = Utils.vec(stats.running_mean)
+    in.running_var = Utils.vec(stats.running_var)
+    return nothing
 end
 
-CRC.@non_differentiable update_instancenorm_state(::Any...)
+CRC.@non_differentiable update_instancenorm_state!(::Any...)
 
 function Base.show(io::IO, l::InstanceNorm)
     print(io, "InstanceNorm($(l.chs)")
@@ -554,18 +554,20 @@ function initialparameters(rng::AbstractRNG, ln::LayerNorm)
     return (;)
 end
 
-function (l::LayerNorm)(x::AbstractArray, ps, st::NamedTuple)
-    x′ = match_eltype(l, ps, st, x)
+function apply(::Type{<:LayerNorm}, l, x::AbstractArray)
+    # XXX: restore `match_eltype` support
+    # x′ = match_eltype(l, ps, st, x)
+    x′ = x
     σ = NNlib.fast_act(l.activation, x′)
     y = layernorm(
         x′,
-        safe_getproperty(ps, Val(:scale)),
-        safe_getproperty(ps, Val(:bias)),
+        safe_getproperty(l, Val(:scale)),
+        safe_getproperty(l, Val(:bias)),
         σ,
         l.dims,
         convert(unwrapped_eltype(x′), l.epsilon),
     )
-    return y, st
+    return y
 end
 
 function Base.show(io::IO, l::LayerNorm)
@@ -781,17 +783,14 @@ parameterlength(l::RMSNorm) = has_affine(l) ? prod(l.normalized_shape) : 0
 
 # specialization on `NT` is important here, else we won't be able to infer the
 # correct eltype of the output.
-function (rms::RMSNorm)(x::AbstractArray{T}, ps, st::NamedTuple) where {T}
+function apply(::Type{<:RMSNorm}, rms, x::AbstractArray{T}) where {T}
     # Don't use `match_eltype` here, since often times the eltypes are intentionally
     # different.
     ϵ = T(rms.epsilon)
     mean_sq = mean(abs2, x; dims=1:length(rms.normalized_shape))
 
     if has_affine(rms)
-        norm_x = @. (x * LuxOps.rsqrt(mean_sq + ϵ)) * ps.scale
-    else
-        norm_x = @. x * LuxOps.rsqrt(mean_sq + ϵ)
+        return @. (x * LuxOps.rsqrt(mean_sq + ϵ)) * ps.scale
     end
-
-    return norm_x, st
+    return @. x * LuxOps.rsqrt(mean_sq + ϵ)
 end