Add jacobian function

src/Flux.jl

@@ -9,7 +9,7 @@ using Lazy: @forward
 
 export Chain, Dense, RNN, LSTM, Dropout, LayerNorm,
   SGD, ADAM, Momentum, Nesterov,
-  param, params, mapleaves
+  param, params, mapleaves, jacobian
 
 using NNlib
 export σ, relu, leakyrelu, elu, swish, softmax

src/utils.jl

@@ -120,3 +120,22 @@ function throttle(f, timeout; leading=true, trailing=false)
     nothing
   end
 end
+
+"""
+    J = jacobian(m,x)
+
+Calculate the output jacobian `J = d/dx m(x)` such that each row `i` of `J` corresponds to the gradient `J[i,:] = ∇ₓ(m(x)[i])`
+"""
+function jacobian(m,x)
+    xp = param(x)
+    y  = m(xp)
+    k  = length(y)
+    n  = length(x)
+    J  = Matrix{eltype(x)}(n,k)
+    for i = 1:k
+        Flux.back!(y[i]) # Populate gradient accumulator
+        J[:,i] = xp.grad
+        xp.grad .*= 0 # Reset gradient accumulator
+    end
+    J'
+end

test/utils.jl

@@ -47,3 +47,12 @@ using Flux: throttle
     @test a == [1, 3]
   end
 end
+
+@testset "Jacobian" begin
+  A = param(randn(2,2))
+  x = randn(2)
+  m(x) = A*x
+  y = m(x)
+  J = jacobian(m,x)
+  @test J ≈ A.data
+end