DifferentiationInterface.jl/DifferentiationInterfaceTest/test/zero_backends.jl at 22bd82eb4ab89ea55f8690166f5847fa1b7a2ab6 · penelopeysm/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
using ADTypes
using DifferentiationInterface
using DifferentiationInterface: AutoZeroForward, AutoZeroReverse
using DifferentiationInterfaceTest
using DifferentiationInterfaceTest: allocfree_scenarios, no_matrices
using Test

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

## Type stability

test_differentiation(
    AutoZeroForward(),
    map(zero, default_scenarios(; include_batchified=false));
    type_stability=:full,
    logging=LOGGING,
)

test_differentiation(
    AutoZeroReverse(),
    map(
        DifferentiationInterfaceTest.same_function,
        default_scenarios(; include_batchified=false),
    );
    correctness=false,
    type_stability=:prepared,
    logging=LOGGING,
)

## Benchmark

data0 = benchmark_differentiation(
    AutoZeroForward(),
    no_matrices(default_scenarios(; include_batchified=false, include_constantified=true));
    logging=LOGGING,
);

data1 = benchmark_differentiation(
    AutoZeroForward(),
    no_matrices(default_scenarios(; include_batchified=false));
    benchmark=:full,
    logging=LOGGING,
    benchmark_seconds=0.05,
    benchmark_aggregation=maximum,
);

struct FakeBackend <: ADTypes.AbstractADType end
ADTypes.mode(::FakeBackend) = ADTypes.ForwardMode()

data2 = benchmark_differentiation(
    FakeBackend(),
    no_matrices(default_scenarios(; include_batchified=false));
    logging=false,
    benchmark_test=false,
);

@testset "Benchmarking DataFrame" begin
    for col in eachcol(data1)
        if eltype(col) <: AbstractFloat
            @test !any(isnan, col)
        end
    end
    for col in eachcol(data2)
        if eltype(col) <: AbstractFloat
            @test all(isnan, col)
        end
    end
end

## Allocations

@testset "Benchmark for zero allocations" begin
    data_allocfree = vcat(
        benchmark_differentiation(
            AutoZeroForward(),
            allocfree_scenarios();
            excluded=[:pullback, :gradient],
            benchmark=:prepared,
            logging=LOGGING,
        ),
        benchmark_differentiation(
            AutoZeroReverse(),
            allocfree_scenarios();
            excluded=[:pushforward, :derivative],
            benchmark=:prepared,
            logging=LOGGING,
        ),
    )
    @testset "$(collect(row[1:4]))" for row in collect(eachrow(data_allocfree))
        @test row[:allocs] == 0
    end
end

test_differentiation(
    AutoZeroForward(),
    allocfree_scenarios();
    correctness=false,
    allocations=:prepared,
    excluded=[:pullback, :gradient, :jacobian],
    logging=LOGGING,
)

test_differentiation(
    AutoZeroReverse(),
    allocfree_scenarios();
    correctness=false,
    allocations=:prepared,
    excluded=[:pushforward, :derivative, :jacobian],
    logging=LOGGING,
)

test_differentiation(
    AutoZeroForward();
    correctness=false,
    allocations=:full,
    skip_allocations=true,
    logging=LOGGING,
)