refactor: remove first(batchnorm_cudnn)

ext/LuxLibTrackercuDNNExt.jl

@@ -6,24 +6,6 @@ using LuxLib: LuxLib
 using Tracker: Tracker, TrackedVector, TrackedArray
 
 # api/batchnorm.jl
-const TR_CUDNN_BN_ARRAY_TYPE = Union{
-    TrackedArray{<:Any, <:Any, <:CuArray{<:Union{Float32, Float64}, 2}},
-    TrackedArray{<:Any, <:Any, <:CuArray{<:Union{Float32, Float64}, 4}},
-    TrackedArray{<:Any, <:Any, <:CuArray{<:Union{Float32, Float64}, 5}}}
-const TR_BNParamType = Union{
-    Nothing, TrackedArray{<:Any, <:Any, <:CuVector{<:Union{Float32, Float64}}},
-    CuVector{<:Union{Float32, Float64}}}
-
-function LuxLib.batchnorm(
-        x::TR_CUDNN_BN_ARRAY_TYPE, scale::TR_BNParamType, bias::TR_BNParamType,
-        running_mean::TR_BNParamType, running_var::TR_BNParamType, training::Val,
-        σ::F=identity, momentum::Real=0.1f0, epsilon::Real=1.0f-5) where {F}
-    rm, rv = LuxLib._get_batchnorm_statistics(x, running_mean, running_var, training)
-    # NOTE: The following returns a tracked tuple so we can't do `first` on it
-    x_ = LuxLib.batchnorm_cudnn(rm, rv, scale, bias, x, momentum, epsilon, training)[1]
-    return LuxLib.fast_activation!!(σ, x_), (; running_mean=rm, running_var=rv)
-end
-
 for RM in (:TrackedVector, :Nothing, :AbstractVector),
     RV in (:TrackedVector, :Nothing, :AbstractVector),
     S in (:TrackedVector, :Nothing, :AbstractVector),

ext/LuxLibcuDNNExt/LuxLibcuDNNExt.jl

@@ -23,7 +23,7 @@ function LuxLib.batchnorm(x::CUDNN_BN_ARRAY_TYPE, scale::BNParamType, bias::BNPa
         running_mean::BNParamType, running_var::BNParamType, training::Val,
         σ::F=identity, momentum::Real=0.1f0, epsilon::Real=1.0f-5) where {F}
     rm, rv = LuxLib._get_batchnorm_statistics(x, running_mean, running_var, training)
-    x_ = first(LuxLib.batchnorm_cudnn(rm, rv, scale, bias, x, momentum, epsilon, training))
+    x_ = LuxLib.batchnorm_cudnn(rm, rv, scale, bias, x, momentum, epsilon, training)[1]
     return LuxLib.fast_activation!!(σ, x_), (; running_mean=rm, running_var=rv)
 end