fix: mooncake lux

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.31.1"
+version = "1.31.2"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -127,7 +127,7 @@ MLUtils = "0.4.4"
 MPI = "0.20.19"
 MacroTools = "0.5.13"
 Markdown = "1.10"
-Mooncake = "0.4.148, 0.5"
+Mooncake = "0.5"
 NCCL = "0.1.2"
 NNlib = "0.9.27"
 Optimisers = "0.4.6"

ext/MooncakeExt/training.jl

@@ -3,13 +3,7 @@ function Lux.Training.compute_gradients_impl(
 ) where {F}
     config = get_config(ad)
     pullback_cache = prepare_pullback_cache(
-        objective_function,
-        ts.model,
-        ts.parameters,
-        ts.states,
-        data;
-        debug_mode=config.debug_mode,
-        silence_debug_messages=config.silence_debug_messages,
+        objective_function, ts.model, ts.parameters, ts.states, data; config
     )
     # evaluate once to get the correct types
     loss, stₙ, stats = objective_function(ts.model, ts.parameters, ts.states, data)

lib/LuxLib/test/others/onehotarray_tests.jl

@@ -1,17 +1,11 @@
-using OneHotArrays, LuxLib, Test, BenchmarkTools
+using OneHotArrays, LuxLib, Test
 
 @testset "Specialized OneHotArrays Dispatch" begin
     x = onehotbatch("aabc", "abcdefghijklmnopqrstuv")
     weight = reshape(collect(Float32, 1:(1024 * 22)), 1024, 22)
 
-    dense_res_time = @belapsed fused_dense_bias_activation(
-        identity, $(weight), $(Array(x)), nothing
-    )
-    onehot_res_time = @belapsed fused_dense_bias_activation(
-        identity, $(weight), $(x), nothing
-    )
-
-    @test onehot_res_time < dense_res_time / 5
     @test fused_dense_bias_activation(identity, weight, x, nothing) ≈
         fused_dense_bias_activation(identity, weight, Array(x), nothing)
+
+    @test LuxLib.Utils.force_3arg_mul!_dispatch(weight, weight, x)
 end