Add jacobian

[oscardssmith]

This uses a slightly different definition of the jacobian than ForwardDiff. For jacobian(f:N×M->P) it will return a result of shape P×N×M while ForwardDiff will return P×(N⋅M)

[codecov-commenter]

Codecov Report

Merging #81 (22a49f6) into main (82096ee) will decrease coverage by 0.30%. The diff coverage is 16.66%.

@@            Coverage Diff             @@
##             main      #81      +/-   ##
==========================================
- Coverage   52.62%   52.31%   -0.31%     
==========================================
  Files          21       21              
  Lines        2172     2177       +5     
==========================================
- Hits         1143     1139       -4     
- Misses       1029     1038       +9     

Impacted Files	Coverage Δ
src/interface.jl	65.00% <16.66%> (-5.91%)	:arrow_down:
src/stage1/forward.jl	74.14% <0.00%> (-2.73%)	:arrow_down:

---

Continue to review full report at Codecov.

Legend - Click here to learn more Δ = absolute <relative> (impact), ø = not affected, ? = missing data Powered by Codecov. Last update 82096ee...22a49f6. Read the comment docs.

[cossio]

What happens for f: N×M -> PxQ? Is the Jacobian P x Q x N x M?

[oscardssmith]

This PR only enables support of jacobians for functions that output a vector.

[mcabbott]

Shouldn't this return a tuple, like gradient, to allow for multiple arguments?

julia> Diffractor.gradient(atan)(1/2)
(0.8,)

julia> Diffractor.gradient(atan)(1, 2)
(0.4, -0.2)

julia> Diffractor.jacobian(x -> x[1:2])([3,4,5])
2×3 Matrix{Int64}:
 1  0  0
 0  1  0

julia> Diffractor.jacobian(*)(rand(2,2), rand(2)) 
ERROR: MethodError: no method matching (::Diffractor.JacobianBack{1, typeof(*)})(...

... and allow jacobian(f, x...) in order to allow do blocks less ugly than this one:

julia> Diffractor.jacobian() do x
         x[1:2]
       end([1,2,3])
2×3 Matrix{Int64}:
 1  0  0
 0  1  0

One argument against this make-more-dimensions story is that few other things work like that. Not just all the other AD packages, but for instance * does not:

julia> rand(2)' * Diffractor.jacobian(x -> x[1:2])(rand(4))
1×4 adjoint(::Vector{Float64}) with eltype Float64:
 0.036297  0.265575  0.0  0.0

julia> rand(2)' * Diffractor.jacobian(x -> x[1:2])(rand(3,4))  # makes a 2×3×4 Array{Float64, 3}
ERROR: MethodError: no method matching *(::Adjoint{Float64, Vector{Float64}}, ::Array{Float64, 3})

If it only accepts arrays of as arguments, then perhaps this should be in its signature, rather than trying & failing. Perhaps only arrays of numbers in fact:

julia> Diffractor.jacobian(x -> x[1])([[1,2], [3,4]])
ERROR: MethodError: no method matching zero(::Type{Vector{Int64}})

julia> Diffractor.jacobian(x -> [x, 2x])(3.0)
ERROR: MethodError: no method matching reshape(::Float64, ::Tuple{Int64})

julia> Diffractor.jacobian(x -> sum(x.a) .+ x.b)((a=[1,2], b=[3,4], c=[5,6])) 
ERROR: MethodError: no method matching size(::NamedTuple{(:a, :b, :c), Tuple{Vector{Int64}, Vector{Int64}, Vector{Int64}}})

Zygote also has various evil tests like this gradient-type-unstable function:

# https://github.com/FluxML/Zygote.jl/blob/dad65a8066669d4c8d7019c28a3dda06ad14c963/test/utils.jl#L68-L71
julia> jacobian(xs -> map(x -> x>0 ? x^3 : 0, xs))([1,2,-3,4,-5])
ERROR: AssertionError: ∂f isa TaylorBundle || ∂f isa TangentBundle{1}