Autograd in Flux

[asyrov]

I guess this may help:

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 …