fix: instance norm gradients with enzyme

Project.toml

@@ -23,7 +23,6 @@ SLEEFPirates = "476501e8-09a2-5ece-8869-fb82de89a1fa"
 StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 UnrolledUtilities = "0fe1646c-419e-43be-ac14-22321958931b"
-VectorizedStatistics = "3b853605-1c98-4422-8364-4bd93ee0529e"
 
 [weakdeps]
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
@@ -77,7 +76,6 @@ Statistics = "1.10"
 Test = "1.10"
 Tracker = "0.2.34"
 UnrolledUtilities = "0.1.2"
-VectorizedStatistics = "0.5.10"
 Zygote = "0.6.70"
 cuDNN = "1.3"
 julia = "1.10"

src/LuxLib.jl

@@ -20,7 +20,6 @@ using StaticArraysCore: StaticArraysCore, StaticArray, StaticVector
 using Statistics: Statistics, mean, var
 using SLEEFPirates: SLEEFPirates
 using UnrolledUtilities: unrolled_any, unrolled_all, unrolled_filter, unrolled_mapreduce
-using VectorizedStatistics: vmean, vvar
 
 @reexport using NNlib
 

src/api/instancenorm.jl

@@ -45,7 +45,7 @@ end
 end
 
 function _test_valid_instancenorm_arguments(::AbstractArray{T, N}) where {T, N}
-    N > 2 || throw(ArgumentError("`ndims(x) = $(N)` must be at least 2."))
+    N > 2 || throw(ArgumentError("`ndims(x) = $(N)` must be at least > 2."))
     return nothing
 end
 

src/impl/fast_ops.jl

@@ -2,17 +2,13 @@
 # VectorizedStatistics.jl, we can will specialize the CPU dispatches to use them.
 fast_mean(x::AbstractArray; dims=:) = fast_mean(internal_operation_mode(x), x; dims)
 fast_mean(opmode, x::AbstractArray; dims=:) = mean(x; dims)
-fast_mean(::LoopedArrayOp, x::AbstractArray; dims=:) = vmean(x; dims, multithreaded=true)
 
 function fast_var(x::AbstractArray; mean=nothing, dims=:, corrected=true)
     return fast_var(internal_operation_mode(x), x; mean, dims, corrected)
 end
 function fast_var(opmode, x::AbstractArray; mean=nothing, dims=:, corrected=true)
     return var(x; mean, dims, corrected)
 end
-function fast_var(::LoopedArrayOp, x::AbstractArray; mean=nothing, dims=:, corrected=true)
-    return vvar(x; mean, dims, corrected, multithreaded=true)
-end
 
 function fast_mean_var(x::AbstractArray; dims=:, corrected=true)
     return fast_mean_var(internal_operation_mode(x), x; dims, corrected)

src/impl/normalization.jl

@@ -17,6 +17,9 @@ function __update_statistics(opmode, rμ, rσ², μ, σ², m1, m2)
     __update_statistics!(opmode, rμ2, rσ²2, rμ, rσ², μ, σ², m1, m2, 1 - m1)
     return rμ2, rσ²2
 end
+
+CRC.@non_differentiable __update_statistics(::Any...)
+
 function __update_statistics!(::LoopedArrayOp, rμ2, rσ²2, rμ, rσ², μ, σ², m1, m2, m3)
     @tturbo for I in indices((rμ2, rσ²2))
         rμ2[I] = m3 * rμ[I] + m1 * μ[I]
@@ -37,7 +40,7 @@ end
     @inbounds rσ²2[I] = m3 * rσ²[I] + m2 * σ²[I]
 end
 
-CRC.@non_differentiable __update_statistics(::Any...)
+EnzymeRules.inactive(::typeof(__update_statistics!), ::Any...) = nothing
 
 function _update_normalization_statistics(
         x::AbstractArray{T, N}, rμ::AbstractArray{<:Number, N},