# ForwardDiff.jacobian of softmax vs analytic Jacobian

**URL:** https://discourse.julialang.org/t/forwarddiff-jacobian-of-softmax-vs-analytic-jacobian/35147
**Category:** Performance
**Tags:** question, differentiation
**Created:** [February 26, 2020, 1:18am UTC](https://discourse.julialang.org/t/forwarddiff-jacobian-of-softmax-vs-analytic-jacobian/35147 "2020-02-26T01:18:24Z")
**Posts on this page:** 1
**Page:** 1

<div class="post-metadata">

### Author: ![zxzkja](https://avatars.discourse-cdn.com/v4/letter/z/71c47a/32.png) [@zxzkja](https://discourse.julialang.org/u/zxzkja)
#### Post date: [February 26, 2020, 1:18am UTC](https://discourse.julialang.org/t/forwarddiff-jacobian-of-softmax-vs-analytic-jacobian/35147/1 "2020-02-26T01:18:24Z")

</div>

I am trying out automatic differentiation on part of my code in place of the analytic derivative, but ForwardDiff.jacobian of softmax seems to take twice as long as the analytic Jacobian. I’m not sure if this is an expected result or I’m doing something to slow down ForwardDiff. Any feedback is appreciated.

Below is a MWE with data representative of my problem. The loop within the functions is also representative (i.e. compute a list of Jacobians for subsets of a larger array). I tried passing an anonymous function to the jacobian, as well as a function that doesn’t compute softmax in-place (both commented out). I know that pushing to an `Array{Any}` is not great, but that’s just for the MWE.

```julia
using ForwardDiff
using BenchmarkTools

# softmax copied from StatsFuns without convert
function softmax2(x::AbstractArray{T}) where {T<:Real}
    n = length(x)
    u = maximum(x)
    s = 0.
    r = similar(x)
    @inbounds for i = 1:n
        s += (r[i] = exp(x[i] - u))
    end
    invs = inv(s)
    @inbounds for i = 1:n
        r[i] *= invs
    end
    r
end
function softmax2!(r::AbstractArray{T}, x::AbstractArray{T}) where {T<:Real}
    n = length(x)
    length(r) == n || throw(DimensionMismatch("Inconsistent array lengths."))
    u = maximum(x)
    s = 0.
    @inbounds for i = 1:n
        s += (r[i] = exp(x[i] - u))
    end
    invs = inv(s)
    @inbounds for i = 1:n
        r[i] *= invs
    end
    r
end

function jac1(x, beta)
    result = Array{Any}(undef, 0)
    for i in 1:10:(size(x)[1] - 9)
        thisx = x[i:(i + 9),:]
        pp!(y::Vector, b::Vector) = softmax2!(y, thisx * b)
        # pp! = (y, b) -> softmax2!(y, thisx * b)
        p = zeros(Float64, 10)
        lambda = ForwardDiff.jacobian(pp!, p, beta)::Array{Float64, 2}
        # pp = b -> softmax2(thisx * b)
        # 𝐩 = pp(beta)
        # lambda = ForwardDiff.jacobian(pp, beta)
        push!(result, lambda)
    end
    return result
end
function jac2(x, beta)
    result = Array{Any}(undef, 0)
    for i in 1:10:(size(x)[1] - 9)
        thisx = x[i:(i + 9),:]
        p = softmax2(thisx * beta)
        lambda = similar(thisx)
        for k in 1:length(beta)
            last_term = 0.
            for s in 1:10
                last_term += thisx[s, k] * p[s]
            end
            for t in 1:10
                lambda[t, k] = p[t] * (thisx[t, k] - last_term)
            end
        end
        push!(result, lambda)
    end
    return result
end

X = [0.398106 0.496961; -0.612026 -0.224875; 0.34112 -1.11714;
     -1.12936 -0.394995; 1.43302 1.54983; 1.9804 -0.743514; -0.367221 -2.33171;
     -1.04413 0.812245; 0.56972 -0.501311; -0.135055 -0.510887;
     2.40162 -1.21536; -0.03924 -0.0225586; 0.689739 0.701239;
     0.0280022 -0.587482; -0.743273 -0.606728; 0.188792 1.09664;
     -1.80496 -0.24751; 1.46555 -0.159902; 0.153253 -0.625778;
     2.17261 0.900435]
β = [0.924499; 0.869371]
X = repeat(X, outer = (5, 2))
β = repeat(β, 2)

display(isapprox(jac1(X, β), jac2(X, β)))

@btime jac1(X, β)
@btime jac2(X, β)

```

Output

```julia
true 8.137 μs (94 allocations: 25.50 KiB)
  4.086 μs (44 allocations: 12.22 KiB)

```