softmax(x) and logsoftmax(x) update their arguments

[dfdx]

Describe the bug

Unlike NNlib versions, softmax(x) and logsoftmax(x) in CUDA.jl modify their arguments in-place.

To reproduce

The simplest way to confirm it is to look at the code, but for the sake of completeness here's the minimal reproducible example:

using CUDA
using NNlib

x = cu(rand(5, 10))
original_x = copy(x)
softmax(x)
@assert x == original_x

Possible fix

It should be easy to fix softmax(), logsoftmax(), ∇softmax() and ∇logsoftmax() by allocating a new array for the result, e.g.:

softmax(x::DenseCuArray{T}; dims=1) where T<:CUDNNFloat =
         softmax!(similar(x), x, dims=dims)

Right now I'm struggling with another issue stopping me from creating a PR, but I plan to come back to it later.

[denizyuret]

I am fixing this and other inconsistencies at #523

[CarloLucibello]

@maleadt would you take a fix to this for backporting on the julia 1.5 compatible branch? This unless #523 is meant to be backported

[maleadt]

@maleadt would you take a fix to this for backporting on the julia 1.5 compatible branch?

Sure. The new CUDNN interface is unlikely to be backported.

[denizyuret]

@maleadt -- I don't think there is much backport work to do for the new cudnn interface, if you let me know what branch, I can check to make sure everything works for 1.5.

[maleadt]

The branch is called compat, I'll make a PR with the CUDNN changes for you and e.g. @DhairyaLGandhi or @CarloLucibello to verify works properly.

[DhairyaLGandhi]

Thanks, that's helpful

[jeremiedb]

@denizyuret I think this issue hasn't been solved at the moment (at least on latest CUDA 2.6.3 release):

using CUDA
using Flux

julia> x = CuArray([-1.0, 0.0, 1.0])
3-element CuArray{Float64, 1}:
 -1.0
  0.0
  1.0

julia> softmax(x)
3-element CuArray{Float64, 1}:
 0.09003057317038046
 0.24472847105479767
 0.6652409557748219

julia> x
3-element CuArray{Float64, 1}:
 0.09003057317038046
 0.24472847105479767
 0.6652409557748219

Could it be reopened so as to keep track of it?

[denizyuret]

I think it is solved in master but not released yet.

…

On Thu, Apr 8, 2021 at 5:54 AM jeremiedb ***@***.***> wrote: @denizyuret <https://github.com/denizyuret> I think this issue hasn't been solved at the moment (at least on latest CUDA 2.6.3 release): using CUDAusing Flux julia> x = CuArray([-1.0, 0.0, 1.0])3-element CuArray{Float64, 1}: -1.0 0.0 1.0 julia> softmax(x)3-element CuArray{Float64, 1}: 0.09003057317038046 0.24472847105479767 0.6652409557748219 julia> x3-element CuArray{Float64, 1}: 0.09003057317038046 0.24472847105479767 0.6652409557748219 Could it be reopened so as to keep track of it? — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <#592 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAN43JVWYXU5VASVIR5WMM3THULE3ANCNFSM4UO33QXQ> .