fix: try fixing broken tests (#1279)

* fix: try fixing AMDGPU kaiming normal

* docs: missing NNlib function

* fix: incorrect package usage

* fix: incorrect package usage

* fix: add device to GenericBroadcastOp

* chore: run the formatter

* test: skip some tests

* test: mark weight initializers tests as broken

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "Lux"
 uuid = "b2108857-7c20-44ae-9111-449ecde12c47"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
-version = "1.11.1"
+version = "1.11.2"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

docs/src/api/NN_Primitives/NNlib.md

@@ -118,6 +118,7 @@ ctc_loss
 ```@docs
 logsumexp
 glu
+NNlib.@disallow_spawns
 ```
 
 !!! tip

ext/LuxReactantExt/patches.jl

@@ -1,4 +1,4 @@
-Utils.vec(x::AnyTracedRArray) = Reactant.TracedUtils.materialize_traced_array(vec(x))
+Utils.vec(x::AnyTracedRArray) = ReactantCore.materialize_traced_array(vec(x))
 
 # XXX: Use PoolDims once EnzymeJAX supports stablehlo.reduce_window adjoint
 Lux.calculate_pool_dims(g::Lux.GlobalPoolMode, ::TracedRArray) = g

lib/LuxLib/Project.toml

@@ -1,7 +1,7 @@
 name = "LuxLib"
 uuid = "82251201-b29d-42c6-8e01-566dec8acb11"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
-version = "1.7.1"
+version = "1.7.2"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"

lib/LuxLib/ext/LuxLibcuDNNExt/LuxLibcuDNNExt.jl

@@ -1,7 +1,7 @@
 module LuxLibcuDNNExt
 
 using LuxLib: LuxLib, Optional, ∂∅, Impl
-using LuxLib.Utils: safe_reshape, safe_vec, unsafe_known, recursive_unthunk
+using LuxLib.Utils: Utils, safe_reshape, safe_vec, unsafe_known, recursive_unthunk
 using CUDA: CUDA, CuArray, CuVector, CU_NULL, DenseCuArray, DenseCuVector
 using ChainRulesCore: ChainRulesCore
 using cuDNN:

lib/LuxLib/ext/LuxLibcuDNNExt/batchnorm.jl

@@ -14,8 +14,8 @@ function Impl.batchnorm_cudnn(::Nothing, ::Nothing, x::DenseCuArray, args...)
 
     y, xμ, xσ⁻² = Impl.batchnorm_cudnn(γ, β, x, args...)
 
-    CUDA.unsafe_free!(γ)
-    CUDA.unsafe_free!(β)
+    Utils.unsafe_free!(γ)
+    Utils.unsafe_free!(β)
 
     return y, xμ, xσ⁻²
 end
@@ -136,10 +136,10 @@ function Impl.∇batchnorm_cudnn(
 
     ∂γ, ∂β, ∂x = Impl.∇batchnorm_cudnn(γ, β, x, ∂y, rμ, rσ², args...)
 
-    CUDA.unsafe_free!(γ)
-    CUDA.unsafe_free!(β)
-    CUDA.unsafe_free!(∂γ)
-    CUDA.unsafe_free!(∂β)
+    Utils.unsafe_free!(γ)
+    Utils.unsafe_free!(β)
+    Utils.unsafe_free!(∂γ)
+    Utils.unsafe_free!(∂β)
 
     return nothing, nothing, ∂x
 end

lib/LuxLib/src/impl/batched_mul.jl

@@ -31,7 +31,7 @@ function batched_matmul(
 end
 
 function batched_matmul(
-    opmode::GPUBroadcastOp{AMDGPUDevice},
+    opmode::Union{GPUBroadcastOp{AMDGPUDevice},GenericBroadcastOp{AMDGPUDevice}},
     x::AbstractArray{<:Complex,3},
     y::AbstractArray{<:Complex,3},
 )

lib/LuxLib/src/traits.jl

@@ -158,7 +158,7 @@ abstract type AbstractInternalArrayOpMode end
 
 abstract type AbstractBroadcastOpMode <: AbstractInternalArrayOpMode end
 
-struct GenericBroadcastOp <: AbstractBroadcastOpMode end
+struct GenericBroadcastOp{dev} <: AbstractBroadcastOpMode end
 struct GPUBroadcastOp{dev} <: AbstractBroadcastOpMode end
 struct LoopedArrayOp <: AbstractInternalArrayOpMode end
 
@@ -192,15 +192,15 @@ Currently supported modes are:
 """
 function internal_operation_mode(xs::Tuple)
     xs = filter(!isnothing, xs)
-    known(Traits.use_generic_broadcasting(xs)) && return GenericBroadcastOp()
-
     dev = get_device_type(xs)
+
+    known(Traits.use_generic_broadcasting(xs)) && return GenericBroadcastOp{dev}()
     dev <: AbstractGPUDevice && return GPUBroadcastOp{dev}()
-    dev <: ReactantDevice && return GenericBroadcastOp()
+    dev <: ReactantDevice && return GenericBroadcastOp{dev}()
 
     # This check needs to be done after the GPU Check
     known(Utils.unrolled_any(!Traits.fast_scalar_indexing, xs)) &&
-        return GenericBroadcastOp()
+        return GenericBroadcastOp{dev}()
     return LoopedArrayOp()
 end
 internal_operation_mode(x::AbstractArray) = internal_operation_mode((x,))

lib/LuxLib/test/others/forwarddiff_tests.jl

@@ -154,6 +154,9 @@
 
                 test_jvp_computation(x -> op(x, kernel_size; stride, pad), x, u, ongpu)
 
+                # NNlib doesn't define ∇meanpool and ∇maxpool for AMDGPU properly
+                mode == "amdgpu" && continue
+
                 test_jvp_computation(
                     x ->
                         only(Zygote.gradient(x -> sum(op(x, kernel_size; stride, pad)), x)),

lib/WeightInitializers/src/initializers.jl

@@ -99,7 +99,7 @@ function kaiming_normal(
 ) where {T<:Number}
     std = T(gain) / sqrt(T(first(Utils.nfan(dims...))))
     x = DeviceAgnostic.randn(rng, T, dims...)
-    x .*= std
+    broadcast!(Base.Fix2(*, std), x, x)
     return x
 end