Latest release (0.7.9) broke softmax gradients with Tracker

[devmotion]

The following example works flawlessly with NNlib 0.7.8 but yields an error with NNlib 0.7.9:

julia> using Tracker, NNlib, LinearAlgebra

julia> Tracker.gradient(x -> dot(x[1:5], softmax(x[6:10]; dims=:)), rand(10)) # same for `softmax(x[6:10]; dims=1)`
ERROR: MethodError: no method matching Float64(::Tracker.TrackedReal{Float64})
Closest candidates are:
  Float64(::Real, ::RoundingMode) where T<:AbstractFloat at rounding.jl:200
  Float64(::T) where T<:Number at boot.jl:716
  Float64(::UInt32) at float.jl:66
  ...
Stacktrace:
 [1] convert(::Type{Float64}, ::Tracker.TrackedReal{Float64}) at ./number.jl:7
 [2] setindex!(::Array{Float64,1}, ::Tracker.TrackedReal{Float64}, ::Int64) at ./array.jl:847
 [3] setindex! at ./array.jl:861 [inlined]
 [4] #374 at /home/david/.julia/packages/Tracker/XcCOb/src/lib/array.jl:105 [inlined]
 [5] back_(::Tracker.Call{Tracker.var"#374#376"{Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},TrackedArray{…,Array{Float64,1}},Tuple{UnitRange{Int6
4}}},Tuple{Tracker.Tracked{Array{Float64,1}},Nothing}}, ::Array{Tracker.TrackedReal{Float64},1}, ::Bool) at /home/david/.julia/packages/Tracker/XcCOb/src/back.jl:35
 [6] back(::Tracker.Tracked{Array{Float64,1}}, ::Array{Tracker.TrackedReal{Float64},1}, ::Bool) at /home/david/.julia/packages/Tracker/XcCOb/src/back.jl:58
 [7] #13 at /home/david/.julia/packages/Tracker/XcCOb/src/back.jl:38 [inlined]
 [8] foreach at ./abstractarray.jl:2010 [inlined]
 [9] back_(::Tracker.Call{Tracker.var"#503#504"{TrackedArray{…,Array{Float64,1}},Array{Tracker.TrackedReal{Float64},1}},Tuple{Tracker.Tracked{Array{Float64,1}},Nothing}},
 ::Tracker.TrackedReal{Float64}, ::Bool) at /home/david/.julia/packages/Tracker/XcCOb/src/back.jl:38
 [10] back(::Tracker.Tracked{Tracker.TrackedReal{Float64}}, ::Int64, ::Bool) at /home/david/.julia/packages/Tracker/XcCOb/src/back.jl:58
 [11] #back!#15 at /home/david/.julia/packages/Tracker/XcCOb/src/back.jl:77 [inlined]
 [12] #back!#32 at /home/david/.julia/packages/Tracker/XcCOb/src/lib/real.jl:16 [inlined]
 [13] back!(::Tracker.TrackedReal{Tracker.TrackedReal{Float64}}) at /home/david/.julia/packages/Tracker/XcCOb/src/lib/real.jl:14
 [14] gradient_(::Function, ::Array{Float64,1}) at /home/david/.julia/packages/Tracker/XcCOb/src/back.jl:4
 [15] #gradient#24 at /home/david/.julia/packages/Tracker/XcCOb/src/back.jl:164 [inlined]
 [16] gradient(::Function, ::Array{Float64,1}) at /home/david/.julia/packages/Tracker/XcCOb/src/back.jl:164
 [17] top-level scope at REPL[32]:1

Due to this problem, the CI tests in DistributionsAD (e.g. https://github.com/TuringLang/DistributionsAD.jl/pull/143/checks?check_run_id=1592162436) and Bijectors fail currently.

[mcabbott]

Presumably #247, cc @norci

[norci]

I have found the root cause:
I used a pre allocated array as output.
Then Tracker does not work.
Zygote also raises Mutating arrays is not supported

function softmax!(out::O, x::T; dims = 1) where {O<:AbstractArray,T<:AbstractArray}
    out .= exp.(x .- maximum(x; dims = dims))
    out ./= sum(out; dims = dims)
end

this works.

    gradient(x -> dot([1.0, 0], softmax(x)), [0.0, 1])

this causes an error

    gradient(x -> dot(x, softmax(x)), [0.0, 1])

@CarloLucibello ,
Do you have any idea?
Or shall we revert softmax to v0.7.8 ?

[CarloLucibello]

Zygote works fine, you can use mutating operations as long as you provide an explicit adjoint,
which we do in ∇softmax.

julia> using Zygote, NNlib, LinearAlgebra  # Zygote v5.17, NNlib v0.7.9

julia> gradient(x -> dot([1.0, 0], softmax(x)), [0.0, 1])
([0.19661193324148185, -0.19661193324148185],)

julia> gradient(x -> dot(x, softmax(x)), [0.0, 1])
([0.07232948812851325, 0.9276705118714867],)

[CarloLucibello]

The problem should be only with Tracker, which must be performing tracing without exploiting the definition of ∇softmax. I'm not very familiar with Tracker, but I think the problem should be fixed there, not here

cc @DhairyaLGandhi

[CarloLucibello]

maybe this will already be fixed by FluxML/Tracker.jl#90

[norci]

gradient(x -> dot(x, softmax(x)), [0.0, 1])

Zygote@0.5.17 works.

But Zygote@0.6.0 raises that error.

[CarloLucibello]

that is fine, we moved NNlib's AD rules out of Zygote in 0.6.0, they are being ported here and will be available in NNlib v0.8

[DhairyaLGandhi]

So current NNlib is broken and shouldn't be used since it's incompatible with zygote 0.5?

[CarloLucibello]

@DhairyaLGandhi ??? current NNlib is perfectly fine and compatible with zygote 0.5

[DhairyaLGandhi]

Sorry it was a question, I mistyped. I edited the question mark in