Make softmax Enzyme-friendly

[mcabbott]

A slight type instability causes this failure with Enzyme:

julia> using NNlib, Enzyme

julia> Enzyme.gradient(Reverse, x -> sum(abs2, softmax(x)), randn(3,4))  # Float64 is fine
3×4 Matrix{Float64}:
 -0.122783   0.0819441   0.197442   -0.123512
 -0.109191  -0.0401521  -0.0706376  -0.0803758
  0.231975  -0.041792   -0.126805    0.203887

julia> gradient(Reverse, x -> sum(abs2, softmax(x)), randn(Float32, 3,4))

Illegal updateAnalysis prev:{[-1]:Pointer, [-1,0]:Float@float} new: {[-1]:Pointer, [-1,0]:Integer, [-1,1]:Integer, [-1,2]:Integer, [-1,3]:Integer, [-1,4]:Integer, [-1,5]:Integer, [-1,6]:Integer, [-1,7]:Integer}
val:   %ifelse_result = select i1 %ifelse_cond419, i64 addrspace(11)* %.0.sroa_cast, i64 addrspace(11)* %.0.sroa_cast3, !dbg !562 origin=  %unbox429 = load i64, i64 addrspace(11)* %ifelse_result, align 4, !dbg !572, !tbaa !568, !alias.scope !570, !noalias !571
MethodInstance for NNlib.var"#softmax!#205"(::Int64, ::typeof(softmax!), ::Matrix{Float32}, ::Matrix{Float32})

Caused by:
Stacktrace:
  [1] ifelse
    @ ./essentials.jl:647
  [2] _broadcast_getindex_evalf
    @ ./broadcast.jl:709
  [3] _broadcast_getindex
    @ ./broadcast.jl:682
  [4] getindex
    @ ./broadcast.jl:636
  [5] macro expansion
    @ ./broadcast.jl:1004
  [6] macro expansion
    @ ./simdloop.jl:77
  [7] copyto!
    @ ./broadcast.jl:1003

This PR fixes it. No tests, I'm a bit surprised Flux's tests didn't notice this.