Open gdalle opened 6 months ago
The documentation says
withjacobian(f, args...) Returns both the value f(args...) and the jacobian as a named tuple.
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:
vec(f(args...))
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 https://github.com/FluxML/Zygote.jl/blob/c0daccded5b9f91d31ceb889e4a97e74dd722a4e/src/lib/grad.jl#L203
The documentation says
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: