Skip to content

withjacobian flattens the output when it is a matrix #1506

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
gdalle opened this issue Mar 7, 2024 · 1 comment · May be fixed by #1568
Open

withjacobian flattens the output when it is a matrix #1506

gdalle opened this issue Mar 7, 2024 · 1 comment · May be fixed by #1568
Labels
bug Something isn't working

Comments

@gdalle
Copy link

gdalle commented Mar 7, 2024

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
@gdalle gdalle mentioned this issue Mar 7, 2024
Merged
@mcabbott mcabbott added the bug Something isn't working label Mar 12, 2024
@gdalle gdalle mentioned this issue Mar 25, 2024
Closed
@gdalle
Copy link
Author

gdalle commented Apr 17, 2024

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)

@mcabbott mcabbott linked a pull request Apr 9, 2025 that will close this issue
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees
No one assigned
Labels
bug Something isn't working
Projects
None yet
Milestone
No milestone
Development

Successfully merging a pull request may close this issue.

Don't return vec(y) in withjacobian mcabbott/Zygote.jl
2 participants
@gdalle @mcabbott