add multi arg jacobian

Author: Roger-luo

src/grad.jl

@@ -34,6 +34,19 @@ function jacobian(fdm, f, x::Vector{T}, D::Int) where {T<:Real}
 end
 jacobian(fdm, f, x::Vector{<:Real}) = jacobian(fdm, f, x, length(f(x)))
 
+function jacobian(fdm, f, x::Real, D::Int)
+    x_vec, vec_to_x = to_vec(x)
+    return jacobian(fdm, f, x_vec, D)
+end
+
+replace_arg(k, xs::Tuple, x) = (xs[1:k-1]..., x, xs[k+1:end]...)
+
+function jacobian(fdm, f, xs...)
+    D = length(f(xs...))
+    N = length(xs)
+    return ntuple(k->jacobian(fdm, x->f(replace_arg(k, xs, x)), xs[k], D), N)
+end
+
 """
     _jvp(fdm, f, x::Vector{<:Real}, ẋ::AbstractVector{<:Real})