Split JacobianOperator

Project.toml

@@ -9,6 +9,9 @@ Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
+[sources]
+Enzyme = {url = "https://github.com/EnzymeAD/Enzyme.jl", rev = "vc/fwdsplit"}
+
 [compat]
 Enzyme = "0.13.50"
 Krylov = "0.10.1"

src/Ariadne.jl

@@ -104,6 +104,111 @@ function mul!(out, J′::Union{Adjoint{<:Any, <:JacobianOperator}, Transpose{<:A
     return nothing
 end
 
+##
+# SplitJacobianOperator - Uses Enzyme's ForwardModeSplit for efficient Newton-Krylov
+##
+
+"""
+    SplitJacobianOperator
+
+Uses Enzyme's ForwardModeSplit to cache the primal computation and avoid
+re-running it on every Krylov iteration. The primal `f!(res, u, p)` is
+computed once during `prepare!` and the tape is reused for all JVP queries.
+"""
+struct SplitJacobianOperator{F, F′, A, P, P′, FwdThunk, DerivThunk} <: AbstractJacobianOperator
+    f::F         # f!(res, u, p)
+    f′::F′       # shadow for f (nothing if guaranteed_const)
+    res::A       # residual buffer (shared with workspace)
+    res′::A      # shadow for res during prepare! (zero tangent; allocated once)
+    u::A         # state
+    u′::A        # shadow for u during prepare! (zero tangent; allocated once)
+    p::P
+    p′::P′       # shadow for p (nothing if p is const)
+    fwd_thunk::FwdThunk
+    deriv_thunk::DerivThunk
+    tape::Base.RefValue{Any}   # tape stored after prepare!
+end
+
+"""
+    SplitJacobianOperator(f::F, res::A, u::A, p::P; assume_p_const=false)
+
+Creates a SplitJacobianOperator using Enzyme's ForwardModeSplit mode.
+Requires Enzyme with ForwardSplitNoPrimal support (PR #3024).
+"""
+function SplitJacobianOperator(f::F, res::A, u::A, p::P; assume_p_const = false) where {F, A, P}
+    f_const = Enzyme.Compiler.guaranteed_const(F)
+    p_const = assume_p_const || Enzyme.Compiler.guaranteed_const(P)
+
+    f′ = f_const ? nothing : Enzyme.make_zero(f)
+    p′ = p_const ? nothing : Enzyme.make_zero(p)
+    res′ = Enzyme.make_zero(res)
+    u′ = Enzyme.make_zero(u)
+
+    FA = f_const ? Const{F} : Duplicated{F}
+    PA = p_const ? Const{P} : Duplicated{P}
+
+    # Check if ForwardSplitNoPrimal is available
+    if !isdefined(Enzyme, :ForwardSplitNoPrimal)
+        error("SplitJacobianOperator requires Enzyme with ForwardSplitNoPrimal support (PR #3024)")
+    end
+
+    fwd_thunk, deriv_thunk = Enzyme.autodiff_thunk(
+        Enzyme.ForwardSplitNoPrimal, FA, Const,
+        Duplicated{A}, Duplicated{A}, PA
+    )
+
+    return SplitJacobianOperator(
+        f, f′, res, res′, u, u′, p, p′,
+        fwd_thunk, deriv_thunk, Ref{Any}(nothing)
+    )
+end
+
+"""
+    prepare!(J::SplitJacobianOperator)
+
+Runs the augmented forward pass to populate `J.res` with `f!(J.res, J.u, J.p)` 
+and stores the tape for subsequent derivative passes.
+"""
+function prepare!(J::SplitJacobianOperator)
+    # res′ and u′ stay zero → tangent for the augmented pass is zero
+    # (we only want the primal result in res and the tape; not a JVP)
+    result = J.fwd_thunk(
+        maybe_duplicated(J.f, J.f′),
+        Duplicated(J.res, J.res′),
+        Duplicated(J.u, J.u′),
+        maybe_duplicated(J.p, J.p′),
+    )
+    # Handle both (tape,) and tape return formats defensively
+    J.tape[] = result isa Tuple ? first(result) : result
+    return nothing
+end
+
+batch_size(::SplitJacobianOperator) = 1
+
+Base.size(J::SplitJacobianOperator) = (length(J.res), length(J.u))
+Base.eltype(J::SplitJacobianOperator) = eltype(J.u)
+Base.length(J::SplitJacobianOperator) = prod(size(J))
+
+function mul!(out, J::SplitJacobianOperator, v)
+    J.deriv_thunk(
+        maybe_duplicated(J.f, J.f′),
+        Duplicated(J.res, reshape(out, size(J.res))),
+        Duplicated(J.u, reshape(v, size(J.u))),
+        maybe_duplicated(J.p, J.p′),
+        J.tape[],
+    )
+    return nothing
+end
+
+LinearAlgebra.adjoint(J::SplitJacobianOperator) = Adjoint(J)
+LinearAlgebra.transpose(J::SplitJacobianOperator) = Transpose(J)
+
+# Note: Adjoint/transpose operations for SplitJacobianOperator would need reverse mode
+# which is not part of ForwardModeSplit. For now, these are not implemented.
+function mul!(out, J′::Union{Adjoint{<:Any, <:SplitJacobianOperator}, Transpose{<:Any, <:SplitJacobianOperator}}, v)
+    error("Adjoint/transpose operations not yet implemented for SplitJacobianOperator")
+end
+
 
 function init_cache(x, ::Val{N}) where {N}
     if !Enzyme.Compiler.guaranteed_const(typeof(x))
@@ -383,6 +488,17 @@ function evaluate!(ws::NewtonKrylovWorkspace)
     return norm(ws.res)
 end
 
+"""
+    Ariadne.evaluate!(ws::NewtonKrylovWorkspace{<:Any,<:Any,<:Any,<:SplitJacobianOperator}) -> norm_res
+
+Specialized evaluate! for SplitJacobianOperator that calls `prepare!` instead of 
+evaluating the function directly. This both populates `ws.res` and caches the tape.
+"""
+function evaluate!(ws::NewtonKrylovWorkspace{<:Any, <:Any, <:Any, <:SplitJacobianOperator})
+    prepare!(ws.J)
+    return norm(ws.res)
+end
+
 ##
 # LineSearches
 ##

test/split_jacobian.jl

@@ -0,0 +1,149 @@
+using Test
+using Ariadne
+import Ariadne: SplitJacobianOperator, prepare!
+using Enzyme, LinearAlgebra
+
+# Skip tests if ForwardSplitNoPrimal is not available
+const FORWARD_SPLIT_AVAILABLE = isdefined(Enzyme, :ForwardSplitNoPrimal)
+
+function F!(res, x, _)
+    res[1] = x[1]^2 + x[2]^2 - 2
+    res[2] = exp(x[1] - 1) + x[2]^2 - 2
+    return nothing
+end
+
+function F(x, p)
+    res = similar(x)
+    F!(res, x, p)
+    return res
+end
+
+@testset "SplitJacobianOperator" begin
+    if !FORWARD_SPLIT_AVAILABLE
+        @test_skip "ForwardSplitNoPrimal not available, skipping SplitJacobianOperator tests"
+        return
+    end
+
+    @testset "constructor" begin
+        x = [3.0, 5.0]
+        res = zeros(2)
+        p = nothing
+
+        J_split = SplitJacobianOperator(F!, res, x, p)
+
+        @test size(J_split) == (2, 2)
+        @test length(J_split) == 4
+        @test eltype(J_split) == Float64
+        @test J_split.tape[] === nothing  # not prepared yet
+    end
+
+    @testset "prepare! and mul!" begin
+        x = [3.0, 5.0]
+        res = zeros(2)
+        p = nothing
+
+        J_split = SplitJacobianOperator(F!, res, x, p)
+
+        # Before prepare!, res should be zeros
+        @test res == [0.0, 0.0]
+
+        # After prepare!, res should contain F!(res, x, p) and tape should be stored
+        prepare!(J_split)
+        expected_res = F(x, p)
+        @test res ≈ expected_res
+        @test J_split.tape[] !== nothing
+
+        # Test JVP computation
+        v = [1.0, 0.0]
+        out = zeros(2)
+        mul!(out, J_split, v)
+
+        # Compare with regular JacobianOperator result
+        J_regular = Ariadne.JacobianOperator(F!, copy(res), x, p)
+        out_regular = zeros(2)
+        mul!(out_regular, J_regular, v)
+
+        @test out ≈ out_regular
+    end
+
+    @testset "consistency with JacobianOperator" begin
+        x = [3.0, 5.0]
+        res = zeros(2)
+        p = nothing
+
+        # Create both operators
+        J_split = SplitJacobianOperator(F!, res, x, p)
+        prepare!(J_split)
+
+        res_regular = copy(res)
+        J_regular = Ariadne.JacobianOperator(F!, res_regular, x, p)
+
+        # Test on multiple directions
+        for v in ([1.0, 0.0], [0.0, 1.0], [1.0, 1.0], rand(2))
+            out_split = zeros(2)
+            out_regular = zeros(2)
+
+            mul!(out_split, J_split, v)
+            mul!(out_regular, J_regular, v)
+
+            @test out_split ≈ out_regular rtol = 1.0e-12
+        end
+    end
+
+    @testset "collect" begin
+        x = [3.0, 5.0]
+        res = zeros(2)
+        p = nothing
+
+        J_split = SplitJacobianOperator(F!, res, x, p)
+        prepare!(J_split)
+
+        # Compare collected matrix with Enzyme jacobian
+        J_matrix_split = collect(J_split)
+        J_matrix_enz = jacobian(Forward, x -> F(x, nothing), x) |> only
+
+        @test J_matrix_split ≈ J_matrix_enz rtol = 1.0e-12
+    end
+
+    @testset "adjoint/transpose error" begin
+        x = [3.0, 5.0]
+        res = zeros(2)
+        p = nothing
+
+        J_split = SplitJacobianOperator(F!, res, x, p)
+        prepare!(J_split)
+
+        v = [1.0, 0.0]
+        out = zeros(2)
+
+        # Adjoint and transpose should throw errors
+        @test_throws ErrorException mul!(out, adjoint(J_split), v)
+        @test_throws ErrorException mul!(out, transpose(J_split), v)
+    end
+
+    @testset "parameter handling" begin
+        # Test with actual parameters
+        function F_param!(res, x, p)
+            res[1] = x[1]^2 + x[2]^2 - p[1]
+            res[2] = exp(x[1] - 1) + x[2]^2 - p[2]
+            return nothing
+        end
+
+        x = [3.0, 5.0]
+        res = zeros(2)
+        p = [2.0, 2.0]
+
+        # Test with assume_p_const=true
+        J_split = SplitJacobianOperator(F_param!, res, x, p; assume_p_const = true)
+        @test J_split.p′ === nothing
+
+        prepare!(J_split)
+        @test res ≈ [3.0^2 + 5.0^2 - 2.0, exp(3.0 - 1) + 5.0^2 - 2.0]
+
+        # Test with assume_p_const=false (if p is not guaranteed const by Enzyme)
+        if !Enzyme.Compiler.guaranteed_const(typeof(p))
+            J_split2 = SplitJacobianOperator(F_param!, copy(res), x, p; assume_p_const = false)
+            @test J_split2.p′ !== nothing
+        end
+    end
+end