Gradient of Flux model wrt to weights

On newer versions of Zygote/Flux, you would have gotten the error message:

julia> grad = Flux.gradient(x -> model(x), some_input) # Gradient evaluated at the inputs
ERROR: output an array, so the gradient is not defined. Perhaps you wanted jacobian.

and indeed Flux.jacobian(x -> model(x), some_input) works. The terminology is generally that the derivative of a scalar function (like a loss) is the “gradient” and the derivative of a vector function is the “jacobian,” so you just need the latter (and to upgrade to the latest version). Of course, if in the end you do have a loss function, you’ll just want to do a gradient rather than explicitly calculating the intermediate jacobian.

1 Like