fix: incorrect dispatch called

LuxDL · Oct 18, 2024 · 0b19476 · 0b19476
1 parent 8d7c497
commit 0b19476
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 10 deletions.
diff --git a/src/impl/batched_mul.jl b/src/impl/batched_mul.jl
@@ -54,15 +54,14 @@ function batched_matmul!(z::AbstractArray{zT, 3}, ::LoopedArrayOp,
     return
 end
 
-function batched_matmul_cpu!(
-        z::AbstractArray{zT, 3}, x::AbstractArray{xT, 3}, y::AbstractArray{yT, 3},
-        α::Number=true, β::Number=false) where {zT, xT, yT}
+function batched_matmul_cpu!(z::AbstractArray{zT, 3}, x::AbstractArray{xT, 3},
+        y::AbstractArray{yT, 3}) where {zT, xT, yT}
     if can_loopvec_args(batchview(z, 1), batchview(x, 1), batchview(y, 1)) &&
        !unsafe_known(explicit_blas_loaded())
-        batched_matmul_loopvec_impl!(z, x, y, α, β)
+        batched_matmul_loopvec_impl!(z, x, y)
         return
     end
-    NNlib.batched_mul!(z, x, y, α, β)
+    NNlib.batched_mul!(z, x, y)
     return
 end
 

diff --git a/test/common_ops/bias_act_tests.jl b/test/common_ops/bias_act_tests.jl
@@ -44,11 +44,9 @@
             @jet bias_act_loss2(act, x, b)
             @jet bias_act_loss3(act, x, b)
 
-            if act !== lisht
-                @test @inferred(Zygote.gradient(bias_act_loss2, act, x, b)) isa Any broken=(T ==
-                                                                                            Float16)
-                @test @inferred(Zygote.gradient(bias_act_loss3, act, x, b)) isa Any broken=(T ==
-                                                                                            Float16)
+            if act !== lisht && T != Float16
+                @test @inferred(Zygote.gradient(bias_act_loss2, act, x, b)) isa Any
+                @test @inferred(Zygote.gradient(bias_act_loss3, act, x, b)) isa Any
             end
 
             @test_gradients(__Fix1(bias_act_loss1, act), x, b; atol, rtol,