fix: materialize softmax (#1332)

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "Reactant"
 uuid = "3c362404-f566-11ee-1572-e11a4b42c853"
 authors = ["William Moses <[email protected]>", "Valentin Churavy <[email protected]>", "Sergio Sánchez Ramírez <[email protected]>", "Paul Berg <[email protected]>", "Avik Pal <[email protected]>", "Mosè Giordano <[email protected]>"]
-version = "0.2.117"
+version = "0.2.118"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

ext/ReactantNNlibExt/Implementations.jl

@@ -7,6 +7,7 @@ for (jlop, hloop) in (
 end
 
 function NNlib.softmax!(out::AnyTracedRArray{T,N}, x::AbstractArray; dims=1) where {T,N}
+    x = T.(Reactant.materialize_traced_array(x))
     max_ = maximum(x; dims)
     diff = exp.(x .- max_)
     @trace if all(isfinite, max_)
@@ -19,6 +20,7 @@ function NNlib.softmax!(out::AnyTracedRArray{T,N}, x::AbstractArray; dims=1) whe
 end
 
 function NNlib.logsoftmax!(out::AnyTracedRArray{T}, x::AbstractArray; dims=1) where {T}
+    x = T.(Reactant.materialize_traced_array(x))
     max_ = maximum(x; dims)
     diff = x .- max_
     @trace if all(isfinite, max_)

test/nn/nnlib.jl

@@ -418,3 +418,12 @@ end
 
     @test @jit(NNlib.upsample_nearest(x_ra, (2, 2))) ≈ NNlib.upsample_nearest(x, (2, 2))
 end
+
+@testset "softmax/logsoftmax reshaped input" begin
+    x = rand(Float32, 3, 4, 5)
+    x_ra = reshape(Reactant.to_rarray(x), 12, 5)
+    x = reshape(x, 12, 5)
+
+    @test @jit(NNlib.softmax(x_ra)) ≈ NNlib.softmax(x)
+    @test @jit(NNlib.logsoftmax(x_ra)) ≈ NNlib.logsoftmax(x)
+end