 # Computing function Jacobian from Zygote.pullback

Zygote documentation on custom adjoints states `dx/dy` (in the notation from the docs) is the Jacobian when `x`, `y` are vectors. I am struggling in relating this to an implementation.

For a function `f : R^N -> R^M` modeled with `f = Chain(Dense(N, H, tanh), Dense(H, M))` with hidden dimension `H`, I intuitively thought `dy/dx` would be an `MxN` matrix, but as we all know intuition often fails so I experimented a bit. Running

``````using Flux, Zygote
N, M, H = 2, 3, 16
g = Chain(Dense(N, H, tanh), Dense(H, M))
x = randn(N, 1)
y, back = Zygote.pullback(g, x)
back([])
``````

gives` DimensionMismatch("matrix A has dimensions (16,3), matrix B has dimensions (0,1)")` so it looks like `dy/dx` is an `MxH` matrix. Okay, cool, I guess this is the first matrix multiplication in the first `back` call. But do I even get an `MxN` matrix if I compute the back pass? No, running `back(ones(size(y)...))` returns an `Nx1` matrix. So how would I compute the “MxN” Jacobian?

Pullback with the basis elements, i.e. [1,0,0,0,0,0], [0,1,0,0,0,0], …

1 Like

You can steal this implementation https://github.com/FluxML/Zygote.jl/issues/98#issuecomment-497739021 or this one https://github.com/FluxML/Zygote.jl/pull/235/files .

1 Like

Thank you for the responses.

So the way to do it is to compute each column in the Jacobian separately? I thought that would be inefficient and that there is a smart way of doing it in a vectorized way out there. I could very well be wrong about this though!

No each row. And if you have a neural network, pulling back the identity matrix will give the Jacobian.

Pulling back the identity matrix will give the Jacobian transpose.

Thank you for your responses, I think I got a good understanding of it now!