No mapreduce in default test scenarios (#282)

* No mapreduce in default test scenarios

* Fix efficiency tests

* Replace randn with rand

* Default alpha

* Fix

* Fix missing ChainRule in 1.6

* Maybe Zygote on 1.6 will finaly shut up

* Zygote HVP uses ForwardDiff

* No CRC second order on 1.6

* Typo

* Version bump

Author: gdalle

DifferentiationInterface/Project.toml

@@ -1,7 +1,7 @@
 name = "DifferentiationInterface"
 uuid = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.5.1"
+version = "0.5.2"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

DifferentiationInterface/ext/DifferentiationInterfaceZygoteExt/DifferentiationInterfaceZygoteExt.jl

@@ -1,9 +1,14 @@
 module DifferentiationInterfaceZygoteExt
 
-using ADTypes: AutoZygote
+using ADTypes: AutoForwardDiff, AutoZygote
 import DifferentiationInterface as DI
 using DifferentiationInterface:
-    NoGradientExtras, NoHessianExtras, NoJacobianExtras, NoPullbackExtras, PullbackExtras
+    HVPExtras,
+    NoGradientExtras,
+    NoHessianExtras,
+    NoJacobianExtras,
+    NoPullbackExtras,
+    PullbackExtras
 using DocStringExtensions
 using Zygote:
     ZygoteRuleConfig, gradient, hessian, jacobian, pullback, withgradient, withjacobian
@@ -88,6 +93,31 @@ function DI.jacobian!(f, jac, backend::AutoZygote, x, extras::NoJacobianExtras)
     return copyto!(jac, DI.jacobian(f, backend, x, extras))
 end
 
+## HVP (with ForwardDiff)
+
+# TODO: find a way to do this without cheating?
+
+struct ZygoteHVPExtras{G,PE} <: HVPExtras
+    ∇f::G
+    pushforward_extras::PE
+end
+
+function DI.prepare_hvp(f, ::AutoZygote, x, v)
+    ∇f(x) = only(gradient(f, x))
+    pushforward_extras = DI.prepare_pushforward(∇f, AutoForwardDiff(), x, v)
+    return ZygoteHVPExtras(∇f, pushforward_extras)
+end
+
+function DI.hvp(f, ::AutoZygote, x, v, extras::ZygoteHVPExtras)
+    @compat (; ∇f, pushforward_extras) = extras
+    return DI.pushforward(∇f, AutoForwardDiff(), x, v, pushforward_extras)
+end
+
+function DI.hvp!(f, p, ::AutoZygote, x, v, extras::ZygoteHVPExtras)
+    @compat (; ∇f, pushforward_extras) = extras
+    return DI.pushforward!(∇f, p, AutoForwardDiff(), x, v, pushforward_extras)
+end
+
 ## Hessian
 
 DI.prepare_hessian(f, ::AutoZygote, x) = NoHessianExtras()

DifferentiationInterface/test/Single/Diffractor.jl

@@ -8,4 +8,6 @@ for backend in [AutoDiffractor()]
     @test !check_hessian(backend; verbose=false)
 end
 
-test_differentiation(AutoDiffractor(); second_order=false, logging=LOGGING);
+test_differentiation(
+    AutoDiffractor(), default_scenarios(; linalg=false); second_order=false, logging=LOGGING
+);

DifferentiationInterface/test/Single/ForwardDiff/efficiency.jl

@@ -3,33 +3,35 @@ using DifferentiationInterface, DifferentiationInterfaceTest
 using ForwardDiff: ForwardDiff
 using SparseConnectivityTracer
 using SparseMatrixColorings
+using Test
+
+function vecexp!(y, x)
+    y .= exp.(x)
+    return nothing
+end
+
+sumexp(x) = sum(exp, x)
 
 @testset verbose = false "Dense efficiency" begin
     # derivative and gradient for `f(x)`
 
     results1 = benchmark_differentiation(
         [AutoForwardDiff()],
-        default_scenarios();
-        outofplace=false,
-        twoarg=false,
-        input_type=Union{Number,AbstractVector},
-        output_type=Number,
-        second_order=false,
-        excluded=[PullbackScenario],
+        [
+            DerivativeScenario(exp; x=1.0, place=:outofplace),
+            GradientScenario(sumexp; x=rand(10), place=:inplace),
+        ];
         logging=LOGGING,
     )
 
     # derivative and jacobian for f!(x, y)
 
     results2 = benchmark_differentiation(
         [AutoForwardDiff()],
-        default_scenarios();
-        outofplace=false,
-        onearg=false,
-        input_type=Union{Number,AbstractVector},
-        output_type=AbstractVector,
-        second_order=false,
-        excluded=[PullbackScenario],
+        [
+            DerivativeScenario(vecexp!; x=1.0, y=zeros(10), place=:inplace),
+            JacobianScenario(vecexp!; x=rand(10), y=zeros(10), place=:inplace),
+        ];
         logging=LOGGING,
     )
 

DifferentiationInterface/test/Single/Zygote.jl

@@ -22,7 +22,14 @@ end
 
 ## Dense backends
 
-test_differentiation(dense_backends; excluded=[SecondDerivativeScenario], logging=LOGGING);
+test_differentiation(
+    AutoChainRules(Zygote.ZygoteRuleConfig());
+    excluded=[SecondDerivativeScenario],
+    second_order=VERSION >= v"1.10",
+    logging=LOGGING,
+);
+
+test_differentiation(AutoZygote(); excluded=[SecondDerivativeScenario], logging=LOGGING);
 
 test_differentiation(
     AutoZygote(),

DifferentiationInterfaceTest/Project.toml

@@ -1,7 +1,7 @@
 name = "DifferentiationInterfaceTest"
 uuid = "a82114a7-5aa3-49a8-9643-716bb13727a3"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.4.1"
+version = "0.4.2"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

DifferentiationInterfaceTest/src/scenarios/component.jl

@@ -1,14 +1,14 @@
 ## Vector to scalar
 
 function comp_to_num(x::ComponentVector)::Number
-    return sum(sin, x.a) + sum(cos, x.b)
+    return sum(sin.(x.a)) + sum(cos.(x.b))
 end
 
 comp_to_num_gradient(x) = ComponentVector(; a=cos.(x.a), b=-sin.(x.b))
 
 function comp_to_num_pushforward(x, dx)
     g = comp_to_num_gradient(x)
-    return dot(g.a, dx.a) + dot(g.b, dx.b)
+    return sum(g.a .* dx.a) + sum(g.b .* dx.b)
 end
 
 function comp_to_num_pullback(x, dy)

DifferentiationInterfaceTest/src/scenarios/default.jl

@@ -140,15 +140,63 @@ end
 
 ## Array to scalar
 
-arr_to_num(x::AbstractArray)::Number = sum(sin, x)
+const DEFAULT_α = 4
+const DEFAULT_β = 6
 
-arr_to_num_gradient(x) = cos.(x)
-arr_to_num_hvp(x, v) = -sin.(x) .* v
+arr_to_num_aux_linalg(x; α, β) = sum(vec(x .^ α) .* transpose(vec(x .^ β)))
+
+function arr_to_num_aux_no_linalg(x; α, β)
+    n = length(x)
+    s = zero(eltype(x))
+    for i in 1:n, j in 1:n
+        s += x[i]^α * x[j]^β
+    end
+    return s
+end
+
+function arr_to_num_aux_gradient(x; α, β)
+    x = Array(x)  # GPU arrays don't like indexing
+    g = similar(x)
+    for k in eachindex(g, x)
+        g[k] = (
+            α * x[k]^(α - 1) * sum(x[j]^β for j in eachindex(x) if j != k) +
+            β * x[k]^(β - 1) * sum(x[i]^α for i in eachindex(x) if i != k) +
+            (α + β) * x[k]^(α + β - 1)
+        )
+    end
+    return g
+end
+
+function arr_to_num_aux_hessian(x; α, β)
+    x = Array(x)  # GPU arrays don't like indexing
+    H = similar(x, length(x), length(x))
+    for k in axes(H, 1), l in axes(H, 2)
+        if k == l
+            H[k, k] = (
+                α * (α - 1) * x[k]^(α - 2) * sum(x[j]^β for j in eachindex(x) if j != k) +
+                β * (β - 1) * x[k]^(β - 2) * sum(x[i]^α for i in eachindex(x) if i != k) +
+                (α + β) * (α + β - 1) * x[k]^(α + β - 2)
+            )
+        else
+            H[k, l] = α * β * (x[k]^(α - 1) * x[l]^(β - 1) + x[k]^(β - 1) * x[l]^(α - 1))
+        end
+    end
+    return H
+end
+
+arr_to_num_linalg(x::AbstractArray)::Number =
+    arr_to_num_aux_linalg(x; α=DEFAULT_α, β=DEFAULT_β)
+arr_to_num_no_linalg(x::AbstractArray)::Number =
+    arr_to_num_aux_no_linalg(x; α=DEFAULT_α, β=DEFAULT_β)
+
+arr_to_num_gradient(x) = arr_to_num_aux_gradient(x; α=DEFAULT_α, β=DEFAULT_β)
+arr_to_num_hessian(x) = arr_to_num_aux_hessian(x; α=DEFAULT_α, β=DEFAULT_β)
 arr_to_num_pushforward(x, dx) = dot(arr_to_num_gradient(x), dx)
 arr_to_num_pullback(x, dy) = arr_to_num_gradient(x) .* dy
-arr_to_num_hessian(x) = Matrix(Diagonal(-sin.(vec(x))))
+arr_to_num_hvp(x, v) = reshape(arr_to_num_hessian(x) * vec(v), size(x))
 
-function arr_to_num_scenarios_onearg(x::AbstractArray)
+function arr_to_num_scenarios_onearg(x::AbstractArray; linalg=true)
+    arr_to_num = linalg ? arr_to_num_linalg : arr_to_num_no_linalg
     # pushforward stays out of place
     scens = AbstractScenario[]
     for place in (:outofplace, :inplace)
@@ -448,24 +496,24 @@ const IMAT = Matrix((1:2) .* transpose(1:3))
 
 Create a vector of [`AbstractScenario`](@ref)s with standard array types.
 """
-function default_scenarios()
+function default_scenarios(; linalg=true)
     return vcat(
         # one argument
-        num_to_num_scenarios_onearg(randn()),
-        num_to_arr_scenarios_onearg(randn(), IVEC),
-        num_to_arr_scenarios_onearg(randn(), IMAT),
-        arr_to_num_scenarios_onearg(randn(6)),
-        arr_to_num_scenarios_onearg(randn(2, 3)),
-        vec_to_vec_scenarios_onearg(randn(6)),
-        vec_to_mat_scenarios_onearg(randn(6)),
-        mat_to_vec_scenarios_onearg(randn(2, 3)),
-        mat_to_mat_scenarios_onearg(randn(2, 3)),
+        num_to_num_scenarios_onearg(rand()),
+        num_to_arr_scenarios_onearg(rand(), IVEC),
+        num_to_arr_scenarios_onearg(rand(), IMAT),
+        arr_to_num_scenarios_onearg(rand(6); linalg),
+        arr_to_num_scenarios_onearg(rand(2, 3); linalg),
+        vec_to_vec_scenarios_onearg(rand(6)),
+        vec_to_mat_scenarios_onearg(rand(6)),
+        mat_to_vec_scenarios_onearg(rand(2, 3)),
+        mat_to_mat_scenarios_onearg(rand(2, 3)),
         # two arguments
-        num_to_arr_scenarios_twoarg(randn(), IVEC),
-        num_to_arr_scenarios_twoarg(randn(), IMAT),
-        vec_to_vec_scenarios_twoarg(randn(6)),
-        vec_to_mat_scenarios_twoarg(randn(6)),
-        mat_to_vec_scenarios_twoarg(randn(2, 3)),
-        mat_to_mat_scenarios_twoarg(randn(2, 3)),
+        num_to_arr_scenarios_twoarg(rand(), IVEC),
+        num_to_arr_scenarios_twoarg(rand(), IMAT),
+        vec_to_vec_scenarios_twoarg(rand(6)),
+        vec_to_mat_scenarios_twoarg(rand(6)),
+        mat_to_vec_scenarios_twoarg(rand(2, 3)),
+        mat_to_mat_scenarios_twoarg(rand(2, 3)),
     )
 end

DifferentiationInterfaceTest/src/scenarios/gpu.jl

@@ -6,23 +6,23 @@ const JLMAT = jl(IMAT)
 
 Create a vector of [`AbstractScenario`](@ref)s with GPU array types from [JLArrays.jl](https://github.com/JuliaGPU/GPUArrays.jl/tree/master/lib/JLArrays).
 """
-function gpu_scenarios()
+function gpu_scenarios(; linalg=true)
     return vcat(
         # one argument
-        num_to_arr_scenarios_onearg(randn(), JLVEC),
-        num_to_arr_scenarios_onearg(randn(), JLMAT),
-        arr_to_num_scenarios_onearg(jl(randn(6))),
-        arr_to_num_scenarios_onearg(jl(randn(2, 3))),
-        vec_to_vec_scenarios_onearg(jl(randn(6))),
-        vec_to_mat_scenarios_onearg(jl(randn(6))),
-        mat_to_vec_scenarios_onearg(jl(randn(2, 3))),
-        mat_to_mat_scenarios_onearg(jl(randn(2, 3))),
+        num_to_arr_scenarios_onearg(rand(), JLVEC),
+        num_to_arr_scenarios_onearg(rand(), JLMAT),
+        arr_to_num_scenarios_onearg(jl(rand(6)); linalg),
+        arr_to_num_scenarios_onearg(jl(rand(2, 3)); linalg),
+        vec_to_vec_scenarios_onearg(jl(rand(6))),
+        vec_to_mat_scenarios_onearg(jl(rand(6))),
+        mat_to_vec_scenarios_onearg(jl(rand(2, 3))),
+        mat_to_mat_scenarios_onearg(jl(rand(2, 3))),
         # two arguments
-        num_to_arr_scenarios_twoarg(randn(), JLVEC),
-        num_to_arr_scenarios_twoarg(randn(), JLMAT),
-        vec_to_vec_scenarios_twoarg(jl(randn(6))),
-        vec_to_mat_scenarios_twoarg(jl(randn(6))),
-        mat_to_vec_scenarios_twoarg(jl(randn(2, 3))),
-        mat_to_mat_scenarios_twoarg(jl(randn(2, 3))),
+        num_to_arr_scenarios_twoarg(rand(), JLVEC),
+        num_to_arr_scenarios_twoarg(rand(), JLMAT),
+        vec_to_vec_scenarios_twoarg(jl(rand(6))),
+        vec_to_mat_scenarios_twoarg(jl(rand(6))),
+        mat_to_vec_scenarios_twoarg(jl(rand(2, 3))),
+        mat_to_mat_scenarios_twoarg(jl(rand(2, 3))),
     )
 end

DifferentiationInterfaceTest/src/scenarios/sparse.jl

@@ -229,11 +229,11 @@ Create a vector of [`AbstractScenario`](@ref)s with sparse array types, focused
 """
 function sparse_scenarios()
     return vcat(
-        sparse_vec_to_vec_scenarios(randn(6)),
-        sparse_vec_to_mat_scenarios(randn(6)),
-        sparse_mat_to_vec_scenarios(randn(2, 3)),
-        sparse_mat_to_mat_scenarios(randn(2, 3)),
-        sparse_vec_to_num_scenarios(randn(6)),
-        sparse_mat_to_num_scenarios(randn(2, 3)),
+        sparse_vec_to_vec_scenarios(rand(6)),
+        sparse_vec_to_mat_scenarios(rand(6)),
+        sparse_mat_to_vec_scenarios(rand(2, 3)),
+        sparse_mat_to_mat_scenarios(rand(2, 3)),
+        sparse_vec_to_num_scenarios(rand(6)),
+        sparse_mat_to_num_scenarios(rand(2, 3)),
     )
 end