Mixed-mode sparse Jacobians

DifferentiationInterface/Project.toml

@@ -1,7 +1,7 @@
 name = "DifferentiationInterface"
 uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.6.21"
+version = "0.6.22"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -61,7 +61,7 @@ ReverseDiff = "1.15.1"
 SparseArrays = "<0.0.1,1"
 SparseConnectivityTracer = "0.5.0,0.6"
 StaticArrays = "1.9.7"
-SparseMatrixColorings = "0.4.5"
+SparseMatrixColorings = "0.4.9"
 Symbolics = "5.27.1, 6"
 Tracker = "0.2.33"
 Zygote = "0.6.69"
@@ -99,4 +99,4 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["ADTypes", "Aqua", "ComponentArrays", "DataFrames", "ExplicitImports", "ForwardDiff", "JET", "JLArrays", "JuliaFormatter", "Pkg", "Random", "SparseArrays", "SparseConnectivityTracer", "SparseMatrixColorings", "StableRNGs", "StaticArrays", "Test"]
+test = ["ADTypes", "Aqua", "ComponentArrays", "DataFrames", "ExplicitImports", "ForwardDiff", "JET", "JLArrays", "JuliaFormatter", "Pkg", "Random", "SparseArrays", "SparseConnectivityTracer", "SparseMatrixColorings", "StableRNGs", "StaticArrays", "Test", "Zygote"]

DifferentiationInterface/docs/src/api.md

@@ -68,6 +68,7 @@ jacobian
 jacobian!
 value_and_jacobian
 value_and_jacobian!
+MixedMode
 ```
 
 ## Second order

DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/DifferentiationInterfaceSparseMatrixColoringsExt.jl

@@ -21,6 +21,8 @@ using DifferentiationInterface:
     PushforwardPerformance,
     inner,
     outer,
+    forward_backend,
+    reverse_backend,
     multibasis,
     pick_batchsize,
     pushforward_performance,
@@ -33,13 +35,32 @@ using SparseMatrixColorings:
     coloring,
     column_colors,
     row_colors,
+    ncolors,
     column_groups,
     row_groups,
     sparsity_pattern,
     decompress!
 import SparseMatrixColorings as SMC
 
+function fy_with_contexts(f, contexts::Vararg{Context,C}) where {C}
+    return (with_contexts(f, contexts...),)
+end
+
+function fy_with_contexts(f!, y, contexts::Vararg{Context,C}) where {C}
+    return (with_contexts(f!, contexts...), y)
+end
+
+abstract type SparseJacobianPrep <: JacobianPrep end
+
+SMC.sparsity_pattern(prep::SparseJacobianPrep) = sparsity_pattern(prep.coloring_result)
+SMC.column_colors(prep::SparseJacobianPrep) = column_colors(prep.coloring_result)
+SMC.column_groups(prep::SparseJacobianPrep) = column_groups(prep.coloring_result)
+SMC.row_colors(prep::SparseJacobianPrep) = row_colors(prep.coloring_result)
+SMC.row_groups(prep::SparseJacobianPrep) = row_groups(prep.coloring_result)
+SMC.ncolors(prep::SparseJacobianPrep) = ncolors(prep.coloring_result)
+
 include("jacobian.jl")
+include("jacobian_mixed.jl")
 include("hessian.jl")
 
 end

DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/hessian.jl

@@ -19,6 +19,7 @@ end
 SMC.sparsity_pattern(prep::SparseHessianPrep) = sparsity_pattern(prep.coloring_result)
 SMC.column_colors(prep::SparseHessianPrep) = column_colors(prep.coloring_result)
 SMC.column_groups(prep::SparseHessianPrep) = column_groups(prep.coloring_result)
+SMC.ncolors(prep::SparseHessianPrep) = ncolors(prep.coloring_result)
 
 ## Hessian, one argument
 

DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian.jl

@@ -1,21 +1,5 @@
-function fy_with_contexts(f, contexts::Vararg{Context,C}) where {C}
-    return (with_contexts(f, contexts...),)
-end
-
-function fy_with_contexts(f!, y, contexts::Vararg{Context,C}) where {C}
-    return (with_contexts(f!, contexts...), y)
-end
-
 ## Preparation
 
-abstract type SparseJacobianPrep <: JacobianPrep end
-
-SMC.sparsity_pattern(prep::SparseJacobianPrep) = sparsity_pattern(prep.coloring_result)
-SMC.column_colors(prep::SparseJacobianPrep) = column_colors(prep.coloring_result)
-SMC.column_groups(prep::SparseJacobianPrep) = column_groups(prep.coloring_result)
-SMC.row_colors(prep::SparseJacobianPrep) = row_colors(prep.coloring_result)
-SMC.row_groups(prep::SparseJacobianPrep) = row_groups(prep.coloring_result)
-
 struct PushforwardSparseJacobianPrep{
     BS<:BatchSizeSettings,
     C<:AbstractColoringResult{:nonsymmetric,:column},

DifferentiationInterface/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian_mixed.jl

@@ -0,0 +1,231 @@
+## Preparation
+
+struct MixedModeSparseJacobianPrep{
+    BSf<:BatchSizeSettings,
+    BSr<:BatchSizeSettings,
+    C<:AbstractColoringResult{:nonsymmetric,:bidirectional},
+    M<:AbstractMatrix{<:Real},
+    Sf<:Vector{<:NTuple},
+    Sr<:Vector{<:NTuple},
+    Rf<:Vector{<:NTuple},
+    Rr<:Vector{<:NTuple},
+    Ef<:PushforwardPrep,
+    Er<:PullbackPrep,
+} <: SparseJacobianPrep
+    batch_size_settings_forward::BSf
+    batch_size_settings_reverse::BSr
+    coloring_result::C
+    compressed_matrix_forward::M
+    compressed_matrix_reverse::M
+    batched_seeds_forward::Sf
+    batched_seeds_reverse::Sr
+    batched_results_forward::Rf
+    batched_results_reverse::Rr
+    pushforward_prep::Ef
+    pullback_prep::Er
+end
+
+function DI.prepare_jacobian(
+    f::F, backend::AutoSparse{<:MixedMode}, x, contexts::Vararg{Context,C}
+) where {F,C}
+    y = f(x, map(unwrap, contexts)...)
+    return _prepare_mixed_sparse_jacobian_aux(y, (f,), backend, x, contexts...)
+end
+
+function DI.prepare_jacobian(
+    f!::F, y, backend::AutoSparse{<:MixedMode}, x, contexts::Vararg{Context,C}
+) where {F,C}
+    return _prepare_mixed_sparse_jacobian_aux(y, (f!, y), backend, x, contexts...)
+end
+
+function _prepare_mixed_sparse_jacobian_aux(
+    y, f_or_f!y::FY, backend::AutoSparse{<:MixedMode}, x, contexts::Vararg{Context,C}
+) where {FY,C}
+    dense_backend = dense_ad(backend)
+    sparsity = jacobian_sparsity(
+        fy_with_contexts(f_or_f!y..., contexts...)..., x, sparsity_detector(backend)
+    )
+    problem = ColoringProblem{:nonsymmetric,:bidirectional}()
+    coloring_result = coloring(
+        sparsity,
+        problem,
+        coloring_algorithm(backend);
+        decompression_eltype=promote_type(eltype(x), eltype(y)),
+    )
+
+    Nf = length(column_groups(coloring_result))
+    Nr = length(row_groups(coloring_result))
+    batch_size_settings_forward = pick_batchsize(forward_backend(dense_backend), Nf)
+    batch_size_settings_reverse = pick_batchsize(reverse_backend(dense_backend), Nr)
+
+    return _prepare_mixed_sparse_jacobian_aux_aux(
+        batch_size_settings_forward,
+        batch_size_settings_reverse,
+        coloring_result,
+        y,
+        f_or_f!y,
+        backend,
+        x,
+        contexts...,
+    )
+end
+
+function _prepare_mixed_sparse_jacobian_aux_aux(
+    batch_size_settings_forward::BatchSizeSettings{Bf},
+    batch_size_settings_reverse::BatchSizeSettings{Br},
+    coloring_result::AbstractColoringResult{:nonsymmetric,:bidirectional},
+    y,
+    f_or_f!y::FY,
+    backend::AutoSparse{<:MixedMode},
+    x,
+    contexts::Vararg{Context,C},
+) where {Bf,Br,FY,C}
+    Nf, Af = batch_size_settings_forward.N, batch_size_settings_forward.A
+    Nr, Ar = batch_size_settings_reverse.N, batch_size_settings_reverse.A
+
+    dense_backend = dense_ad(backend)
+
+    groups_forward = column_groups(coloring_result)
+    groups_reverse = row_groups(coloring_result)
+
+    seeds_forward = [
+        multibasis(backend, x, eachindex(x)[group]) for group in groups_forward
+    ]
+    seeds_reverse = [
+        multibasis(backend, y, eachindex(y)[group]) for group in groups_reverse
+    ]
+
+    compressed_matrix_forward = stack(_ -> vec(similar(y)), groups_forward; dims=2)
+    compressed_matrix_reverse = stack(_ -> vec(similar(x)), groups_reverse; dims=1)
+
+    batched_seeds_forward = [
+        ntuple(b -> seeds_forward[1 + ((a - 1) * Bf + (b - 1)) % Nf], Val(Bf)) for a in 1:Af
+    ]
+    batched_seeds_reverse = [
+        ntuple(b -> seeds_reverse[1 + ((a - 1) * Br + (b - 1)) % Nr], Val(Br)) for a in 1:Ar
+    ]
+
+    batched_results_forward = [
+        ntuple(b -> similar(y), Val(Bf)) for _ in batched_seeds_forward
+    ]
+    batched_results_reverse = [
+        ntuple(b -> similar(x), Val(Br)) for _ in batched_seeds_reverse
+    ]
+
+    pushforward_prep = prepare_pushforward(
+        f_or_f!y...,
+        forward_backend(dense_backend),
+        x,
+        batched_seeds_forward[1],
+        contexts...,
+    )
+    pullback_prep = prepare_pullback(
+        f_or_f!y...,
+        reverse_backend(dense_backend),
+        x,
+        batched_seeds_reverse[1],
+        contexts...,
+    )
+
+    return MixedModeSparseJacobianPrep(
+        batch_size_settings_forward,
+        batch_size_settings_reverse,
+        coloring_result,
+        compressed_matrix_forward,
+        compressed_matrix_reverse,
+        batched_seeds_forward,
+        batched_seeds_reverse,
+        batched_results_forward,
+        batched_results_reverse,
+        pushforward_prep,
+        pullback_prep,
+    )
+end
+
+## Common auxiliaries
+
+function _sparse_jacobian_aux!(
+    f_or_f!y::FY,
+    jac,
+    prep::MixedModeSparseJacobianPrep{<:BatchSizeSettings{Bf},<:BatchSizeSettings{Br}},
+    backend::AutoSparse,
+    x,
+    contexts::Vararg{Context,C},
+) where {FY,Bf,Br,C}
+    (;
+        batch_size_settings_forward,
+        batch_size_settings_reverse,
+        coloring_result,
+        compressed_matrix_forward,
+        compressed_matrix_reverse,
+        batched_seeds_forward,
+        batched_seeds_reverse,
+        batched_results_forward,
+        batched_results_reverse,
+        pushforward_prep,
+        pullback_prep,
+    ) = prep
+
+    dense_backend = dense_ad(backend)
+    Nf = batch_size_settings_forward.N
+    Nr = batch_size_settings_reverse.N
+
+    pushforward_prep_same = prepare_pushforward_same_point(
+        f_or_f!y...,
+        pushforward_prep,
+        forward_backend(dense_backend),
+        x,
+        batched_seeds_forward[1],
+        contexts...,
+    )
+    pullback_prep_same = prepare_pullback_same_point(
+        f_or_f!y...,
+        pullback_prep,
+        reverse_backend(dense_backend),
+        x,
+        batched_seeds_reverse[1],
+        contexts...,
+    )
+
+    for a in eachindex(batched_seeds_forward, batched_results_forward)
+        pushforward!(
+            f_or_f!y...,
+            batched_results_forward[a],
+            pushforward_prep_same,
+            forward_backend(dense_backend),
+            x,
+            batched_seeds_forward[a],
+            contexts...,
+        )
+
+        for b in eachindex(batched_results_forward[a])
+            copyto!(
+                view(compressed_matrix_forward, :, 1 + ((a - 1) * Bf + (b - 1)) % Nf),
+                vec(batched_results_forward[a][b]),
+            )
+        end
+    end
+
+    for a in eachindex(batched_seeds_reverse, batched_results_reverse)
+        pullback!(
+            f_or_f!y...,
+            batched_results_reverse[a],
+            pullback_prep_same,
+            reverse_backend(dense_backend),
+            x,
+            batched_seeds_reverse[a],
+            contexts...,
+        )
+
+        for b in eachindex(batched_results_reverse[a])
+            copyto!(
+                view(compressed_matrix_reverse, 1 + ((a - 1) * Br + (b - 1)) % Nr, :),
+                vec(batched_results_reverse[a][b]),
+            )
+        end
+    end
+
+    decompress!(jac, compressed_matrix_reverse, compressed_matrix_forward, coloring_result)
+
+    return jac
+end

DifferentiationInterface/src/DifferentiationInterface.jl

@@ -33,6 +33,7 @@ using LinearAlgebra: dot
 
 include("compat.jl")
 
+include("first_order/mixed_mode.jl")
 include("second_order/second_order.jl")
 
 include("utils/prep.jl")
@@ -66,7 +67,7 @@ include("misc/zero_backends.jl")
 ## Exported
 
 export Context, Constant, Cache
-export SecondOrder
+export MixedMode, SecondOrder
 
 export value_and_pushforward!, value_and_pushforward
 export value_and_pullback!, value_and_pullback

DifferentiationInterface/src/first_order/mixed_mode.jl

@@ -0,0 +1,27 @@
+"""
+    MixedMode
+
+Combination of a forward and a reverse mode backend for mixed-mode Jacobian computation.
+
+!!! danger
+    `MixedMode` backends only support [`jacobian`](@ref) and its variants.
+
+# Constructor
+
+    MixedMode(forward_backend, reverse_backend)
+"""
+struct MixedMode{F<:AbstractADType,R<:AbstractADType} <: AbstractADType
+    forward::F
+    reverse::R
+    function MixedMode(forward::AbstractADType, reverse::AbstractADType)
+        @assert pushforward_performance(forward) isa PushforwardFast
+        @assert pullback_performance(reverse) isa PullbackFast
+        return new{typeof(forward),typeof(reverse)}(forward, reverse)
+    end
+end
+
+forward_backend(m::MixedMode) = m.forward
+reverse_backend(m::MixedMode) = m.reverse
+
+struct ForwardAndReverseMode <: ADTypes.AbstractMode end
+ADTypes.mode(::MixedMode) = ForwardAndReverseMode()

DifferentiationInterface/src/utils/batchsize.jl

@@ -52,6 +52,12 @@ function pick_batchsize(backend::AbstractADType, x_or_N::Union{AbstractArray,Int
                 "You should select the batch size for the dense backend of $backend"
             ),
         )
+    elseif backend isa MixedMode
+        throw(
+            ArgumentError(
+                "You should select the batch size for the forward or reverse backend of $backend",
+            ),
+        )
     else
         return BatchSizeSettings(backend, x_or_N)
     end