DifferentiationInterface.jl/DifferentiationInterface/test/Back/Enzyme/test.jl at 6de1319c038b218cc8a3f4387d9f308e4e4073d5 · JuliaDiff/DifferentiationInterface.jl · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
using Pkg
Pkg.add("Enzyme")

using ADTypes: ADTypes
using DifferentiationInterface, DifferentiationInterfaceTest
import DifferentiationInterfaceTest as DIT
using Enzyme: Enzyme
using LinearAlgebra
using StaticArrays
using Test

using ExplicitImports
check_no_implicit_imports(DifferentiationInterface)

LOGGING = get(ENV, "CI", "false") == "false"

function remove_matrix_inputs(scens::Vector{<:Scenario})  # TODO: remove
    if VERSION < v"1.11"
        return scens
    else
        # for https://github.com/EnzymeAD/Enzyme.jl/issues/2071
        return filter(s -> s.x isa Union{Number,AbstractVector}, scens)
    end
end

backends = [
    AutoEnzyme(; mode=nothing),
    AutoEnzyme(; mode=Enzyme.Forward),
    AutoEnzyme(; mode=Enzyme.Reverse),
    AutoEnzyme(; mode=nothing, function_annotation=Enzyme.Const),
]

duplicated_backends = [
    AutoEnzyme(; mode=Enzyme.Forward, function_annotation=Enzyme.Duplicated),
    AutoEnzyme(; mode=Enzyme.Reverse, function_annotation=Enzyme.Duplicated),
]

@testset "Checks" begin
    @testset "Check $(typeof(backend))" for backend in backends
        @test check_available(backend)
        @test check_inplace(backend)
    end
end;

@testset "First order" begin
    test_differentiation(
        backends, default_scenarios(); excluded=SECOND_ORDER, logging=LOGGING
    )

    test_differentiation(
        backends[1:3],
        default_scenarios(; include_normal=false, include_constantified=true);
        excluded=SECOND_ORDER,
        logging=LOGGING,
    )

    test_differentiation(
        backends[2],
        default_scenarios(;
            include_normal=false,
            include_cachified=true,
            include_constantorcachified=true,
            use_tuples=true,
        );
        excluded=SECOND_ORDER,
        logging=LOGGING,
    )

    test_differentiation(
        duplicated_backends,
        default_scenarios(; include_normal=false, include_closurified=true);
        excluded=SECOND_ORDER,
        logging=LOGGING,
    )
end

#=
# TODO: reactivate type stability tests

test_differentiation(
    AutoEnzyme(; mode=Enzyme.Forward),  # TODO: add more
    default_scenarios(; include_batchified=false);
    correctness=false,
    type_stability=:prepared,
    excluded=SECOND_ORDER,
    logging=LOGGING,
);
=#

@testset "Second order" begin
    test_differentiation(
        [
            AutoEnzyme(),
            SecondOrder(
                AutoEnzyme(; mode=Enzyme.Reverse), AutoEnzyme(; mode=Enzyme.Forward)
            ),
        ],
        remove_matrix_inputs(default_scenarios(; include_constantified=true));
        excluded=FIRST_ORDER,
        logging=LOGGING,
    )

    test_differentiation(
        AutoEnzyme(; mode=Enzyme.Forward);
        excluded=vcat(FIRST_ORDER, [:hessian, :hvp]),
        logging=LOGGING,
    )
end

@testset "Sparse" begin
    test_differentiation(
        MyAutoSparse.(AutoEnzyme(; function_annotation=Enzyme.Const)),
        if VERSION < v"1.11"
            sparse_scenarios()
        else
            filter(sparse_scenarios()) do s
                # for https://github.com/EnzymeAD/Enzyme.jl/issues/2168
                (s.x isa AbstractVector) &&
                    (s.f != DIT.sumdiffcube) &&
                    (s.f != DIT.sumdiffcube_mat)
            end
        end;
        sparsity=true,
        logging=LOGGING,
    )
end

@testset "Static" begin
    filtered_static_scenarios = filter(static_scenarios()) do s
        DIT.operator_place(s) == :out && DIT.function_place(s) == :out
    end

    test_differentiation(
        [AutoEnzyme(; mode=Enzyme.Forward), AutoEnzyme(; mode=Enzyme.Reverse)],
        filtered_static_scenarios;
        excluded=SECOND_ORDER,
        logging=LOGGING,
    )
end

@testset "Coverage" begin
    # ConstantOrCache without cache
    f_nocontext(x, p) = x
    @test I == DifferentiationInterface.jacobian(
        f_nocontext, AutoEnzyme(; mode=Enzyme.Forward), rand(10), ConstantOrCache(nothing)
    )
    @test I == DifferentiationInterface.jacobian(
        f_nocontext, AutoEnzyme(; mode=Enzyme.Reverse), rand(10), ConstantOrCache(nothing)
    )
end