DifferentiationInterface.jl/DifferentiationInterface/test/Core/SimpleFiniteDiff/test.jl at ad1c9c88ae7905fd7789b7d7d8761f9b631bf314 · 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
151
152
153
154
155
156
157
using DifferentiationInterface, DifferentiationInterfaceTest
using DifferentiationInterface:
    AutoSimpleFiniteDiff,
    AutoForwardFromPrimitive,
    AutoReverseFromPrimitive,
    DenseSparsityDetector
using SparseMatrixColorings
using JLArrays, StaticArrays
using Test

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

backends = [ #
    AutoSimpleFiniteDiff(; chunksize=5),
    AutoForwardFromPrimitive(AutoSimpleFiniteDiff(; chunksize=4)),
    AutoReverseFromPrimitive(AutoSimpleFiniteDiff(; chunksize=4)),
]

second_order_backends = [ #
    SecondOrder(
        AutoForwardFromPrimitive(AutoSimpleFiniteDiff(; chunksize=5)),
        AutoReverseFromPrimitive(AutoSimpleFiniteDiff(; chunksize=4)),
    ),
    SecondOrder(
        AutoReverseFromPrimitive(AutoSimpleFiniteDiff(; chunksize=5)),
        AutoForwardFromPrimitive(AutoSimpleFiniteDiff(; chunksize=4)),
    ),
]

second_order_hvp_backends = [ #
    SecondOrder(
        AutoReverseFromPrimitive(AutoSimpleFiniteDiff(); inplace=false),
        AutoForwardFromPrimitive(AutoSimpleFiniteDiff()),
    ),
    SecondOrder(
        AutoForwardFromPrimitive(AutoSimpleFiniteDiff(); inplace=false),
        AutoReverseFromPrimitive(AutoSimpleFiniteDiff();),
    ),
    SecondOrder(
        AutoForwardFromPrimitive(AutoSimpleFiniteDiff(); inplace=false),
        AutoForwardFromPrimitive(AutoSimpleFiniteDiff();),
    ),
    SecondOrder(
        AutoReverseFromPrimitive(AutoSimpleFiniteDiff(); inplace=false),
        AutoReverseFromPrimitive(AutoSimpleFiniteDiff();),
    ),
]

adaptive_backends = [ #
    AutoSimpleFiniteDiff(),
    AutoReverseFromPrimitive(AutoSimpleFiniteDiff()),
    SecondOrder(AutoSimpleFiniteDiff(), AutoReverseFromPrimitive(AutoSimpleFiniteDiff())),
    SecondOrder(AutoReverseFromPrimitive(AutoSimpleFiniteDiff()), AutoSimpleFiniteDiff()),
]

for backend in vcat(backends, second_order_backends)
    @test check_available(backend)
    @test check_inplace(backend)
end

## Dense scenarios

@testset "Dense" begin
    test_differentiation(
        vcat(backends, second_order_backends),
        default_scenarios(; include_constantified=true, include_smaller=true);
        logging=LOGGING,
    )

    test_differentiation(
        second_order_hvp_backends,
        default_scenarios(; include_constantorcachified=true);
        excluded=vcat(FIRST_ORDER, :hessian, :second_derivative),
        logging=LOGGING,
    )

    test_differentiation(backends, complex_scenarios(); logging=LOGGING)
end

@testset "Sparse" begin
    test_differentiation(
        MyAutoSparse.(adaptive_backends),
        default_scenarios(; include_constantified=true);
        logging=LOGGING,
    )

    test_differentiation(
        MyAutoSparse.(
            vcat(adaptive_backends, MixedMode(adaptive_backends[1], adaptive_backends[2]))
        ),
        sparse_scenarios(;
            include_constantified=true,
            include_cachified=true,
            include_constantorcachified=true,
            use_tuples=true,
        );
        sparsity=true,
        logging=LOGGING,
    )

    @testset "Complex numbers" begin
        test_differentiation(
            AutoSparse.(
                vcat(
                    adaptive_backends, MixedMode(adaptive_backends[1], adaptive_backends[2])
                );
                sparsity_detector=DenseSparsityDetector(AutoSimpleFiniteDiff(); atol=1e-5),
                coloring_algorithm=GreedyColoringAlgorithm(),
            ),
            complex_sparse_scenarios();
            logging=LOGGING,
        )
    end

    @testset "SparseMatrixColorings access" begin
        jac_for_prep = prepare_jacobian(copy, MyAutoSparse(adaptive_backends[1]), rand(10))
        jac_rev_prep = prepare_jacobian(copy, MyAutoSparse(adaptive_backends[2]), rand(10))
        hess_prep = prepare_hessian(
            x -> sum(abs2, x), MyAutoSparse(adaptive_backends[1]), rand(10)
        )

        @test all(==(1), column_colors(jac_for_prep))
        @test all(==(1), row_colors(jac_rev_prep))
        @test all(==(1), column_colors(hess_prep))
        @test ncolors(jac_for_prep) == 1
        @test ncolors(hess_prep) == 1
        @test only(column_groups(jac_for_prep)) == 1:10
        @test only(row_groups(jac_rev_prep)) == 1:10
        @test only(column_groups(hess_prep)) == 1:10
    end

    @testset "Empty colors for mixed mode" begin # issue 857
        backend = MyAutoSparse(MixedMode(adaptive_backends[1], adaptive_backends[2]))
        @test jacobian(copyto!, zeros(10), backend, ones(10)) isa AbstractMatrix
    end
end

@testset "Misc" begin
    @test_throws ArgumentError DifferentiationInterface.overloaded_input(
        pushforward, sum, AutoSimpleFiniteDiff(), 1, (1, 2)
    )
    @test_throws ArgumentError DifferentiationInterface.overloaded_input(
        pushforward, copyto!, [1.0], AutoSimpleFiniteDiff(), [1.0], ([1.0], [1.0])
    )
end

@testset "Weird arrays" begin
    test_differentiation(
        [
            AutoSimpleFiniteDiff(),
            AutoForwardFromPrimitive(AutoSimpleFiniteDiff()),
            AutoReverseFromPrimitive(AutoSimpleFiniteDiff()),
        ],
        vcat(static_scenarios(), gpu_scenarios());
        logging=LOGGING,
    )
end;