# Computing function Jacobian from Zygote.pullback

**URL:** <https://discourse.julialang.org/t/computing-function-jacobian-from-zygote-pullback/42979>\
**Category:** Machine Learning\
**Tags:** question, flux, zygote\
**Created:** [July 13, 2020, 12:32pm UTC](https://discourse.julialang.org/t/computing-function-jacobian-from-zygote-pullback/42979 "2020-07-13T12:32:30Z")\
**Posts on this page:** 7\
**Page:** 1

<div class="post-metadata">

**Author:** ![SebastianCallh](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sebastiancallh/32/16102_2.png) [@SebastianCallh](https://discourse.julialang.org/u/SebastianCallh)\
**Post date:** [July 13, 2020, 12:32pm UTC](https://discourse.julialang.org/t/computing-function-jacobian-from-zygote-pullback/42979/1 "2020-07-13T12:32:30Z")

</div>

Zygote [documentation on custom adjoints](https://fluxml.ai/Zygote.jl/dev/adjoints/#Pullbacks-1) 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

```julia
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?

I am clearly not thinking about this the right way. Could someone give me some pointers to this problem please? Reading the documentation has unfortunately not made it “click” for me.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [July 13, 2020, 12:59pm UTC](https://discourse.julialang.org/t/computing-function-jacobian-from-zygote-pullback/42979/2 "2020-07-13T12:59:08Z")

</div>

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

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [July 13, 2020, 1:08pm UTC](https://discourse.julialang.org/t/computing-function-jacobian-from-zygote-pullback/42979/3 "2020-07-13T13:08:42Z")

</div>

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

---

<div class="post-metadata">

**Author:** ![SebastianCallh](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sebastiancallh/32/16102_2.png) [@SebastianCallh](https://discourse.julialang.org/u/SebastianCallh)\
**Post date:** [July 13, 2020, 3:08pm UTC](https://discourse.julialang.org/t/computing-function-jacobian-from-zygote-pullback/42979/4 "2020-07-13T15:08:00Z")

</div>

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!

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [July 13, 2020, 3:25pm UTC](https://discourse.julialang.org/t/computing-function-jacobian-from-zygote-pullback/42979/5 "2020-07-13T15:25:44Z")

</div>

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

---

<div class="post-metadata">

**Author:** ![mohamed82008](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mohamed82008/32/18171_2.png) [@mohamed82008](https://discourse.julialang.org/u/mohamed82008)\
**Post date:** [July 13, 2020, 11:10pm UTC](https://discourse.julialang.org/t/computing-function-jacobian-from-zygote-pullback/42979/6 "2020-07-13T23:10:31Z")

</div>

> [@SebastianCallh](#):
>
> 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!

Pulling back the identity matrix will give the Jacobian transpose.

---

<div class="post-metadata">

**Author:** ![SebastianCallh](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sebastiancallh/32/16102_2.png) [@SebastianCallh](https://discourse.julialang.org/u/SebastianCallh)\
**Post date:** [July 14, 2020, 5:56am UTC](https://discourse.julialang.org/t/computing-function-jacobian-from-zygote-pullback/42979/7 "2020-07-14T05:56:03Z")

</div>

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