Add batch/vectorized finite difference Jacobian evaluation

Implements feature requested in #210: allows computing the full Jacobian
in a single batched function call instead of N sequential calls. This is
useful for GPU-parallelized functions that can evaluate multiple inputs
simultaneously.

Adds `batch=true` keyword to `finite_difference_jacobian` and
`finite_difference_jacobian!`. When enabled, `f` receives a matrix where
each column is an input point and returns a matrix of outputs. Supports
forward, central, and complex step methods.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: ChrisRackauckas

src/jacobians.jl

@@ -186,6 +186,167 @@ function _make_Ji(::AbstractArray, xtype, dx, color_i, nrows, ncols)
     size(Ji) != (nrows, ncols) ? reshape(Ji, (nrows, ncols)) : Ji #branch when size(dx) == (1,) => size(Ji) == (1,) while size(J) == (1,1)
 end
 
+"""
+    _finite_difference_jacobian_batch(f, x, fdtype, returntype, f_in; relstep, absstep, dir)
+
+Internal function implementing vectorized/batched finite difference Jacobian computation.
+
+When `batch=true` is passed to `finite_difference_jacobian`, this function is called instead
+of the standard column-by-column approach. The function `f` is expected to accept a matrix
+where each column is an input point, and return a matrix where each column is the
+corresponding output. This allows GPU-parallelized or otherwise vectorized functions to
+evaluate all perturbations in a single call.
+
+For forward differences, a single call to `f` is made with `n+1` columns (base point + `n`
+perturbations) if `f_in` is not provided, or `n` columns if `f_in` is provided.
+For central differences, `2n` columns are used (forward and backward perturbations).
+For complex step, `n` columns are used with complex perturbations.
+"""
+function _finite_difference_jacobian_batch(f, x, fdtype, returntype, f_in;
+        relstep, absstep, dir)
+    fdtype isa Type && (fdtype = fdtype())
+    n = length(x)
+    vecx = _vec(x)
+
+    if fdtype == Val(:forward)
+        epsilons = [compute_epsilon(Val(:forward), vecx[i], relstep, absstep, dir) for i in 1:n]
+
+        if f_in isa Nothing
+            # Include x as the first column so we only call f once
+            X = repeat(vecx, 1, n + 1)
+            for i in 1:n
+                X[i, i + 1] += epsilons[i]
+            end
+            FX = f(X)
+            fx_col = @view FX[:, 1]
+            J = similar(FX, size(FX, 1), n)
+            for i in 1:n
+                @. J[:, i] = (FX[:, i + 1] - fx_col) / epsilons[i]
+            end
+        else
+            X = repeat(vecx, 1, n)
+            for i in 1:n
+                X[i, i] += epsilons[i]
+            end
+            FX = f(X)
+            vfx = _vec(f_in)
+            J = similar(FX, size(FX, 1), n)
+            for i in 1:n
+                @. J[:, i] = (FX[:, i] - vfx) / epsilons[i]
+            end
+        end
+        return J
+
+    elseif fdtype == Val(:central)
+        epsilons = [compute_epsilon(Val(:central), vecx[i], relstep, absstep, dir) for i in 1:n]
+
+        # Build matrix with 2n columns: [x+eps1*e1, x-eps1*e1, x+eps2*e2, x-eps2*e2, ...]
+        X = repeat(vecx, 1, 2n)
+        for i in 1:n
+            X[i, 2i - 1] += epsilons[i]
+            X[i, 2i] -= epsilons[i]
+        end
+        FX = f(X)
+        J = similar(FX, size(FX, 1), n)
+        for i in 1:n
+            @. J[:, i] = (FX[:, 2i - 1] - FX[:, 2i]) / (2 * epsilons[i])
+        end
+        return J
+
+    elseif fdtype == Val(:complex) && returntype <: Real
+        epsilon = eps(eltype(x))
+
+        # Build complex matrix with n columns
+        X = repeat(complex.(vecx), 1, n)
+        for i in 1:n
+            X[i, i] += im * epsilon
+        end
+        FX = f(X)
+        J = similar(FX, real(eltype(FX)), size(FX, 1), n)
+        for i in 1:n
+            @. J[:, i] = imag(FX[:, i]) / epsilon
+        end
+        return J
+    else
+        fdtype_error(returntype)
+    end
+end
+
+"""
+    _finite_difference_jacobian_batch!(J, f, x, fdtype, returntype, f_in; relstep, absstep, dir)
+
+Internal in-place function implementing vectorized/batched finite difference Jacobian computation.
+
+When `batch=true` is passed to `finite_difference_jacobian!`, this function is called instead
+of the standard column-by-column approach. The function `f` is expected to accept two matrix
+arguments `f(FX, X)` where `X` has columns of input points and `FX` is filled with the
+corresponding outputs.
+"""
+function _finite_difference_jacobian_batch!(J, f, x, fdtype, returntype, f_in;
+        relstep, absstep, dir)
+    fdtype isa Type && (fdtype = fdtype())
+    m, n = size(J)
+    vecx = _vec(x)
+
+    if fdtype == Val(:forward)
+        epsilons = [compute_epsilon(Val(:forward), vecx[i], relstep, absstep, dir) for i in 1:n]
+
+        if f_in isa Nothing
+            # n+1 columns: base point + n perturbations
+            X = repeat(vecx, 1, n + 1)
+            for i in 1:n
+                X[i, i + 1] += epsilons[i]
+            end
+            FX = similar(x, m, n + 1)
+            f(FX, X)
+            for i in 1:n
+                @. J[:, i] = (FX[:, i + 1] - FX[:, 1]) / epsilons[i]
+            end
+        else
+            X = repeat(vecx, 1, n)
+            for i in 1:n
+                X[i, i] += epsilons[i]
+            end
+            FX = similar(x, m, n)
+            f(FX, X)
+            vfx = _vec(f_in)
+            for i in 1:n
+                @. J[:, i] = (FX[:, i] - vfx) / epsilons[i]
+            end
+        end
+
+    elseif fdtype == Val(:central)
+        epsilons = [compute_epsilon(Val(:central), vecx[i], relstep, absstep, dir) for i in 1:n]
+
+        X = repeat(vecx, 1, 2n)
+        for i in 1:n
+            X[i, 2i - 1] += epsilons[i]
+            X[i, 2i] -= epsilons[i]
+        end
+        FX = similar(x, m, 2n)
+        f(FX, X)
+        for i in 1:n
+            @. J[:, i] = (FX[:, 2i - 1] - FX[:, 2i]) / (2 * epsilons[i])
+        end
+
+    elseif fdtype == Val(:complex) && returntype <: Real
+        epsilon = eps(eltype(x))
+
+        X = repeat(complex.(vecx), 1, n)
+        for i in 1:n
+            X[i, i] += im * epsilon
+        end
+        FX = similar(X, Complex{eltype(x)}, m, n)
+        f(FX, X)
+        for i in 1:n
+            @. J[:, i] = imag(FX[:, i]) / epsilon
+        end
+    else
+        fdtype_error(returntype)
+    end
+    nothing
+end
+
 """
     FiniteDiff.finite_difference_jacobian(
         f,
@@ -246,7 +407,12 @@ function finite_difference_jacobian(f, x,
         colorvec = 1:length(x),
         sparsity = nothing,
         jac_prototype = nothing,
-        dir = true)
+        dir = true,
+        batch = false)
+    if batch
+        return _finite_difference_jacobian_batch(f, x, fdtype, returntype, f_in;
+            relstep = relstep, absstep = absstep, dir = dir)
+    end
     if f_in isa Nothing
         fx = f(x)
     else
@@ -452,7 +618,13 @@ function finite_difference_jacobian!(J,
         relstep = default_relstep(fdtype, eltype(x)),
         absstep = relstep,
         colorvec = 1:length(x),
-        sparsity = ArrayInterface.has_sparsestruct(J) ? J : nothing)
+        sparsity = ArrayInterface.has_sparsestruct(J) ? J : nothing,
+        batch = false)
+    if batch
+        _finite_difference_jacobian_batch!(J, f, x, fdtype, returntype, f_in;
+            relstep = relstep, absstep = absstep, dir = true)
+        return nothing
+    end
     if f_in isa Nothing && fdtype == Val(:forward)
         if size(J, 1) == length(x)
             fx = zero(x)

test/finitedifftests.jl

@@ -612,3 +612,81 @@ end
     @test FiniteDiff.finite_difference_hessian(f, x1, FiniteDiff.HessianCache(x1)) == Diagonal(2*ones(4))
     @test FiniteDiff.finite_difference_hessian(f, x1, FiniteDiff.HessianCache(x2)) == Diagonal(2*ones(4))
 end
+
+# Batched Jacobian tests (issue #210)
+@time @testset "Batched Jacobian tests" begin
+    # Out-of-place batched function: f(X) where X is n×k, returns m×k
+    function oopf_scalar(x)
+        [(x[1] + 3) * (x[2]^3 - 7) + 18,
+            sin(x[2] * exp(x[1]) - 1)]
+    end
+    function oopf_batch(X::AbstractMatrix)
+        hcat([oopf_scalar(X[:, j]) for j in 1:size(X, 2)]...)
+    end
+    # Also handle single vector input for the batch function
+    oopf_batch(x::AbstractVector) = oopf_scalar(x)
+
+    # In-place batched function: f(FX, X) where X is n×k, FX is m×k
+    function iipf_batch(FX::AbstractMatrix, X::AbstractMatrix)
+        for j in 1:size(X, 2)
+            FX[1, j] = (X[1, j] + 3) * (X[2, j]^3 - 7) + 18
+            FX[2, j] = sin(X[2, j] * exp(X[1, j]) - 1)
+        end
+    end
+
+    x = [1.5, 0.7]
+    J_ref = [[-7 + x[2]^3 3 * (3 + x[1]) * x[2]^2];
+             [exp(x[1]) * x[2] * cos(1 - exp(x[1]) * x[2]) exp(x[1]) * cos(1 - exp(x[1]) * x[2])]]
+
+    @testset "Out-of-place batch" begin
+        J_fwd = FiniteDiff.finite_difference_jacobian(oopf_batch, x, Val{:forward}; batch=true)
+        @test err_func(J_fwd, J_ref) < 1e-6
+
+        # With f_in provided
+        f_in = oopf_scalar(x)
+        J_fwd2 = FiniteDiff.finite_difference_jacobian(oopf_batch, x, Val{:forward}, eltype(x), f_in; batch=true)
+        @test err_func(J_fwd2, J_ref) < 1e-6
+
+        J_cen = FiniteDiff.finite_difference_jacobian(oopf_batch, x, Val{:central}; batch=true)
+        @test err_func(J_cen, J_ref) < 1e-8
+
+        J_cpx = FiniteDiff.finite_difference_jacobian(oopf_batch, x, Val{:complex}; batch=true)
+        @test err_func(J_cpx, J_ref) < 1e-14
+    end
+
+    @testset "In-place batch" begin
+        J = zero(J_ref)
+        FiniteDiff.finite_difference_jacobian!(J, iipf_batch, x, Val{:forward}; batch=true)
+        @test err_func(J, J_ref) < 1e-6
+
+        # With f_in provided
+        f_in = oopf_scalar(x)
+        J .= 0
+        FiniteDiff.finite_difference_jacobian!(J, iipf_batch, x, Val{:forward}, eltype(x), f_in; batch=true)
+        @test err_func(J, J_ref) < 1e-6
+
+        J .= 0
+        FiniteDiff.finite_difference_jacobian!(J, iipf_batch, x, Val{:central}; batch=true)
+        @test err_func(J, J_ref) < 1e-8
+
+        J .= 0
+        FiniteDiff.finite_difference_jacobian!(J, iipf_batch, x, Val{:complex}; batch=true)
+        @test err_func(J, J_ref) < 1e-14
+    end
+
+    @testset "Batch matches non-batch" begin
+        # Test on a larger function to make sure batch and non-batch agree
+        f_oop(x) = [x[1]^2 + x[2]*x[3], sin(x[1]) + x[3]^2, x[1]*x[2]*x[3]]
+        function f_batch(X::AbstractMatrix)
+            hcat([f_oop(X[:, j]) for j in 1:size(X, 2)]...)
+        end
+        f_batch(x::AbstractVector) = f_oop(x)
+
+        x3 = [2.0, 3.0, 1.5]
+        for fdtype in (Val{:forward}, Val{:central}, Val{:complex})
+            J_std = FiniteDiff.finite_difference_jacobian(f_oop, x3, fdtype)
+            J_bat = FiniteDiff.finite_difference_jacobian(f_batch, x3, fdtype; batch=true)
+            @test J_std ≈ J_bat
+        end
+    end
+end