withjacobian flattens the output when it is a matrix #1506
Description
The documentation says
withjacobian(f, args...)Returns both the value f(args...) and the jacobian as a named tuple.
But it seems it returns vec(f(args...)) instead. That is coherent with the Jacobian presented as a matrix (instead of a higher-order tensor) but still surprising:
julia> using Zygote
julia> (; val, grad) = Zygote.withjacobian(identity, rand(2, 3));
julia> val
6-element Vector{Float64}:
0.3708410915090188
0.5970333190918194
0.6445738753265045
0.26354708948211447
0.6709067822518385
0.9436576156842928
julia> only(grad)
6×6 Matrix{Float64}:
1.0 0.0 0.0 0.0 0.0 0.0
0.0 1.0 0.0 0.0 0.0 0.0
0.0 0.0 1.0 0.0 0.0 0.0
0.0 0.0 0.0 1.0 0.0 0.0
0.0 0.0 0.0 0.0 1.0 0.0
0.0 0.0 0.0 0.0 0.0 1.0