FluxML / Zygote.jl

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

I think it must be due to this line

Zygote.jl/src/lib/grad.jl

Line 203 in c0daccd

y, back = pullback(_jvec∘f, pars)

`withjacobian` flattens the output when it is a matrix