search

EnzymeAD / Enzyme.jl

Julia bindings for the Enzyme automatic differentiator
https://enzyme.mit.edu
MIT License
428 stars 59 forks source link

Reverse mode `jacobian` is wrong for matrix input #1391

Open gdalle opened 3 months ago

gdalle commented 3 months ago

This seems related to the wrong jacobian shape of #1332 but it turns out even a reshape is not enough:

julia> using Enzyme

julia> x = float.(reshape(1:6, 2, 3))
2×3 Matrix{Float64}:
 1.0  3.0  5.0
 2.0  4.0  6.0

julia> f(x) = fill(sum(abs2, x), 6);

julia> jac_forward = Enzyme.jacobian(Enzyme.Forward, f, x)
6×6 Matrix{Float64}:
 2.0  4.0  6.0  8.0  10.0  12.0
 2.0  4.0  6.0  8.0  10.0  12.0
 2.0  4.0  6.0  8.0  10.0  12.0
 2.0  4.0  6.0  8.0  10.0  12.0
 2.0  4.0  6.0  8.0  10.0  12.0
 2.0  4.0  6.0  8.0  10.0  12.0

julia> jac_reverse = reshape(Enzyme.jacobian(Enzyme.Reverse, f, x, Val(6), Val(1)), 6, 6)
6×6 Matrix{Float64}:
  2.0   2.0   2.0   4.0   4.0   4.0
  6.0   6.0   6.0   8.0   8.0   8.0
 10.0  10.0  10.0  12.0  12.0  12.0
  2.0   2.0   2.0   4.0   4.0   4.0
  6.0   6.0   6.0   8.0   8.0   8.0
 10.0  10.0  10.0  12.0  12.0  12.0