Add first sparsity support (#124)

Author: gdalle

Project.toml

@@ -81,11 +81,12 @@ JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["ADTypes", "Aqua", "Chairmarks", "DataFrames", "Diffractor", "Documenter", "Enzyme", "FastDifferentiation", "FiniteDiff", "FiniteDifferences", "ForwardDiff", "JET", "JuliaFormatter", "Pkg", "PolyesterForwardDiff", "ReverseDiff", "SparseDiffTools", "Symbolics", "Test", "Tracker", "Zygote"]
+test = ["ADTypes", "Aqua", "Chairmarks", "DataFrames", "Diffractor", "Documenter", "Enzyme", "FastDifferentiation", "FiniteDiff", "FiniteDifferences", "ForwardDiff", "JET", "JuliaFormatter", "Pkg", "PolyesterForwardDiff", "ReverseDiff", "SparseArrays", "SparseDiffTools", "Symbolics", "Test", "Tracker", "Zygote"]

ext/DifferentiationInterfaceSparseDiffToolsExt/DifferentiationInterfaceSparseDiffToolsExt.jl

@@ -2,7 +2,26 @@ module DifferentiationInterfaceSparseDiffToolsExt
 
 using ADTypes
 import DifferentiationInterface as DI
-using SparseDiffTools: JacPrototypeSparsityDetection, SymbolicsSparsityDetection
+using SparseDiffTools:
+    AutoSparseEnzyme,
+    JacPrototypeSparsityDetection,
+    SymbolicsSparsityDetection,
+    sparse_jacobian,
+    sparse_jacobian!,
+    sparse_jacobian_cache
 using Symbolics: Symbolics
 
+# used with @eval to avoid Unions and thus ambiguities
+SPARSE_BACKENDS = [
+    AutoSparseEnzyme,
+    AutoSparseFiniteDiff,
+    AutoSparseForwardDiff,
+    AutoSparsePolyesterForwardDiff,
+    AutoSparseReverseDiff,
+    AutoSparseZygote,
+]
+
+include("allocating.jl")
+include("mutating.jl")
+
 end

ext/DifferentiationInterfaceSparseDiffToolsExt/allocating.jl

@@ -0,0 +1,30 @@
+for AutoSparse in SPARSE_BACKENDS
+    @eval begin
+
+        ## Jacobian
+
+        function DI.prepare_jacobian(f, backend::$AutoSparse, x::AbstractArray)
+            return sparse_jacobian_cache(
+                backend, SymbolicsSparsityDetection(), f, x; fx=f(x)
+            )
+        end
+
+        function DI.value_and_jacobian!!(f, jac, backend::$AutoSparse, x, cache)
+            sparse_jacobian!(jac, backend, cache, f, x)
+            return f(x), jac
+        end
+
+        function DI.jacobian!!(f, jac, backend::$AutoSparse, x, cache)
+            sparse_jacobian!(jac, backend, cache, f, x)
+            return jac
+        end
+
+        function DI.value_and_jacobian(f, backend::$AutoSparse, x, cache)
+            return f(x), sparse_jacobian(backend, cache, f, x)
+        end
+
+        function DI.jacobian(f, backend::$AutoSparse, x, cache)
+            return sparse_jacobian(backend, cache, f, x)
+        end
+    end
+end

ext/DifferentiationInterfaceSparseDiffToolsExt/mutating.jl

@@ -0,0 +1,19 @@
+for AutoSparse in SPARSE_BACKENDS
+    @eval begin
+        ## Jacobian
+
+        function DI.prepare_jacobian(
+            f!, backend::$AutoSparse, y::AbstractArray, x::AbstractArray
+        )
+            return sparse_jacobian_cache(
+                backend, SymbolicsSparsityDetection(), f!, similar(y), x
+            )
+        end
+
+        function DI.value_and_jacobian!!(f!, y, jac, backend::$AutoSparse, x, cache)
+            sparse_jacobian!(jac, backend, cache, f!, y, x)
+            f!(y, x)
+            return y, jac
+        end
+    end
+end

ext/DifferentiationInterfaceTapirExt/DifferentiationInterfaceTapirExt.jl

@@ -7,23 +7,34 @@ using Tapir: CoDual, build_rrule, value_and_pullback!!, zero_codual
 
 DI.supports_mutation(::AutoTapir) = DI.MutationNotSupported()
 
+function zero_sametype!!(x_target, x::Number)
+    return zero(x)
+end
+
+function zero_sametype!!(x_target, x::AbstractArray)
+    x_sametype = convert(typeof(x), x_target)
+    x_sametype .= zero(eltype(x))
+    return x_sametype
+end
+
 function DI.value_and_pullback(f, ::AutoTapir, x, dy, rrule)
     y = f(x)
     dy_righttype = convert(typeof(y), dy)
     _, (_, dx) = value_and_pullback!!(rrule, dy_righttype, f, x)
     return y, dx
 end
 
-for op in [
-    :pushforward,
-    :pullback,
-    :derivative,
-    :gradient,
-    :jacobian,
-    :second_derivative,
-    :hvp,
-    :hessian,
-]
+function DI.value_and_pullback!!(f, dx, ::AutoTapir, x, dy, rrule)
+    y = f(x)
+    dy_righttype = convert(typeof(y), dy)
+    dx_righttype = zero_sametype!!(dx, x)
+    new_y, (_, new_dx) = value_and_pullback!!(
+        rrule, dy_righttype, zero_codual(f), CoDual(x, dx_righttype)
+    )
+    return new_y, new_dx
+end
+
+for op in [:pushforward, :pullback, :derivative, :gradient, :jacobian]
     prep_op = Symbol(:prepare_, op)
     @eval function DI.$prep_op(f, backend::AutoTapir, x)
         return build_rrule(f, x)

lib/DifferentiationInterfaceTest/Project.toml

@@ -12,6 +12,7 @@ DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 

lib/DifferentiationInterfaceTest/src/DifferentiationInterfaceTest.jl

@@ -29,6 +29,7 @@ import DifferentiationInterface as DI
 using JET: @test_call, @test_opt
 using JLArrays: jl
 using LinearAlgebra: Diagonal, dot
+using SparseArrays: SparseArrays, nnz, SparseMatrixCSC
 using StaticArrays: SVector, SMatrix
 using Test: @testset, @test
 
@@ -47,6 +48,7 @@ include("utils/filter.jl")
 include("tests/correctness.jl")
 include("tests/type_stability.jl")
 include("tests/call_count.jl")
+include("tests/sparsity.jl")
 include("tests/benchmark.jl")
 include("tests/test.jl")
 

lib/DifferentiationInterfaceTest/src/tests/sparsity.jl

@@ -0,0 +1,52 @@
+function test_sparsity(ba::AbstractADType, ::typeof(jacobian), scen::Scenario{false};)
+    (; f, x, y, ref) = new_scen = deepcopy(scen)
+    extras = prepare_jacobian(f, ba, x)
+    jac_true = if ref isa AbstractADType
+        jacobian(f, ref, x)
+    else
+        ref.jacobian(x)
+    end
+
+    _, jac1 = value_and_jacobian(f, ba, x, extras)
+    _, jac2 = value_and_jacobian!!(f, mysimilar(jac_true), ba, x, extras)
+
+    jac3 = jacobian(f, ba, x, extras)
+    jac4 = jacobian!!(f, mysimilar(jac_true), ba, x, extras)
+
+    @testset "Sparse type" begin
+        @test jac1 isa SparseMatrixCSC
+        @test jac2 isa SparseMatrixCSC
+        @test jac3 isa SparseMatrixCSC
+        @test jac4 isa SparseMatrixCSC
+    end
+    @testset "Sparsity pattern" begin
+        @test nnz(jac1) < length(jac_true)
+        @test nnz(jac2) < length(jac_true)
+        @test nnz(jac3) < length(jac_true)
+        @test nnz(jac4) < length(jac_true)
+    end
+    return nothing
+end
+
+function test_sparsity(ba::AbstractADType, ::typeof(jacobian), scen::Scenario{true};)
+    (; f, x, y, dy, ref) = new_scen = deepcopy(scen)
+    f! = f
+    extras = prepare_jacobian(f!, ba, y, x)
+    jac_shape = Matrix{eltype(y)}(undef, length(y), length(x))
+    jac_true = if ref isa AbstractADType
+        last(value_and_jacobian!!(f!, mysimilar(y), mysimilar(jac_shape), ref, x))
+    else
+        ref.jacobian(x)
+    end
+
+    y10 = mysimilar(y)
+    _, jac1 = value_and_jacobian!!(f!, y10, mysimilar(jac_true), ba, x, extras)
+
+    @testset "Sparse type" begin
+        @test jac1 isa SparseMatrixCSC
+    end
+    @testset "Sparsity pattern" begin
+        @test nnz(jac1) < length(jac_true)
+    end
+    return nothing
+end

lib/DifferentiationInterfaceTest/src/tests/test.jl

@@ -15,6 +15,7 @@ Testing:
 - `correctness=true`: whether to compare the differentiation results with the theoretical values specified in each scenario. If a backend object like `correctness=AutoForwardDiff()` is passed instead of a boolean, the results will be compared using that reference backend as the ground truth. 
 - `call_count=false`: whether to check that the function is called the right number of times
 - `type_stability=false`: whether to check type stability with JET.jl (thanks to `@test_opt`)
+- `sparsity`: whether to check sparsity of the jacobian / hessian
 - `detailed=false`: whether to print a detailed or condensed test log
 
 Filtering:
@@ -42,6 +43,7 @@ function test_differentiation(
     correctness::Union{Bool,AbstractADType}=true,
     type_stability::Bool=false,
     call_count::Bool=false,
+    sparsity::Bool=false,
     detailed=false,
     # filtering
     input_type::Type=Any,
@@ -68,7 +70,8 @@ function test_differentiation(
         "Differentiation tests -" *
         (correctness != false ? " correctness" : "") *
         (call_count ? " calls" : "") *
-        (type_stability ? " types" : "")
+        (type_stability ? " types" : "") *
+        (sparsity ? " sparsity" : "")
 
     @testset verbose = true "$title" begin
         @testset verbose = detailed "$(backend_string(backend))" for backend in backends
@@ -77,7 +80,7 @@ function test_differentiation(
                     compatible(backend, op, scen)
                 end
                     logging &&
-                        @info "Testing: $(backend_string(backend)) - $op - $(string(scen))"
+                        @info "$title: $(backend_string(backend)) - $op - $(string(scen))"
                     correctness != false && @testset "Correctness" begin
                         test_correctness(backend, op, scen; isapprox, atol, rtol)
                     end
@@ -87,6 +90,9 @@ function test_differentiation(
                     type_stability && @testset "Type stability" begin
                         test_jet(backend, op, scen)
                     end
+                    sparsity && @testset "Sparsity" begin
+                        test_sparsity(backend, op, scen)
+                    end
                 end
             end
         end

src/DifferentiationInterface.jl

@@ -53,7 +53,6 @@ struct AutoTapir <: AbstractReverseMode end
 include("second_order.jl")
 include("traits.jl")
 include("utils.jl")
-include("prepare.jl")
 
 include("pushforward.jl")
 include("pullback.jl")