DifferentiationInterface.jl/DifferentiationInterface/test/Back/ForwardDiff/test.jl at 00324ef1d9f0fe1c762263b04ccefae97caeb79e · 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
using Pkg
Pkg.add("ForwardDiff")

using ADTypes: ADTypes
using ComponentArrays: ComponentArrays
using DifferentiationInterface, DifferentiationInterfaceTest
import DifferentiationInterface as DI
import DifferentiationInterfaceTest as DIT
using ForwardDiff: ForwardDiff
using StaticArrays: StaticArrays, @SVector
using JLArrays: JLArrays
using Test

using ExplicitImports
check_no_implicit_imports(DifferentiationInterface)

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

struct MyTag end

backends = [
    AutoForwardDiff(),
    AutoForwardDiff(; chunksize=5),
    AutoForwardDiff(; tag=ForwardDiff.Tag(MyTag(), Float64)),
]

for backend in backends
    @test check_available(backend)
    @test check_inplace(backend)
end

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

    test_differentiation(
        AutoForwardDiff(),
        default_scenarios(;
            include_normal=false,
            include_batchified=false,
            include_cachified=true,
            include_constantorcachified=true,
            use_tuples=true,
            include_smaller=true,
        );
        logging=LOGGING,
    )

    test_differentiation(
        AutoForwardDiff();
        correctness=false,
        type_stability=safetypestab(:prepared),
        logging=LOGGING,
    )

    test_differentiation(
        AutoForwardDiff(; chunksize=5);
        correctness=false,
        type_stability=safetypestab(:full),
        excluded=[:hessian],
        logging=LOGGING,
    )
end

@testset "Sparse" begin
    test_differentiation(
        MyAutoSparse(AutoForwardDiff()), default_scenarios(); logging=LOGGING
    )

    test_differentiation(
        MyAutoSparse(AutoForwardDiff()),
        sparse_scenarios(; include_constantified=true);
        sparsity=true,
        logging=LOGGING,
    )
end

@testset "Weird" begin
    test_differentiation(AutoForwardDiff(), component_scenarios(); logging=LOGGING)
    test_differentiation(AutoForwardDiff(), static_scenarios(); logging=LOGGING)
    test_differentiation(
        DI.AutoForwardFromPrimitive(AutoForwardDiff()), gpu_scenarios(); logging=LOGGING
    )

    @testset "Batch size" begin
        @test DI.pick_batchsize(AutoForwardDiff(), rand(7)) isa DI.BatchSizeSettings{7}
        @test DI.pick_batchsize(AutoForwardDiff(; chunksize=5), rand(7)) isa
            DI.BatchSizeSettings{5}
        @test (@inferred DI.pick_batchsize(AutoForwardDiff(), @SVector(rand(7)))) isa
            DI.BatchSizeSettings{7}
        @test (@inferred DI.pick_batchsize(
            AutoForwardDiff(; chunksize=5), @SVector(rand(7))
        )) isa DI.BatchSizeSettings{5}
    end
end

@testset verbose = true "Overloaded inputs" begin
    backend = AutoForwardDiff()
    sparse_backend = MyAutoSparse(AutoForwardDiff())

    # Derivative
    x = 1.0
    y = [1.0, 1.0]
    @test DI.overloaded_input_type(prepare_derivative(copy, backend, x)) ==
        ForwardDiff.Dual{ForwardDiff.Tag{typeof(copy),Float64},Float64,1}
    @test DI.overloaded_input_type(prepare_derivative(copyto!, y, backend, x)) ==
        Vector{ForwardDiff.Dual{ForwardDiff.Tag{typeof(copyto!),Float64},Float64,1}}

    # Gradient
    x = [1.0, 1.0]
    @test DI.overloaded_input_type(prepare_gradient(sum, backend, x)) ==
        Vector{ForwardDiff.Dual{ForwardDiff.Tag{typeof(sum),Float64},Float64,2}}

    # Jacobian
    x = [1.0, 0.0, 0.0]
    @test DI.overloaded_input_type(prepare_jacobian(copy, backend, x)) ==
        ForwardDiff.Dual{ForwardDiff.Tag{typeof(copy),Float64},Float64,3}
    @test DI.overloaded_input_type(prepare_jacobian(copyto!, similar(x), backend, x)) ==
        Vector{ForwardDiff.Dual{ForwardDiff.Tag{typeof(copyto!),Float64},Float64,3}}
    @test DI.overloaded_input_type(
        prepare_jacobian(copyto!, similar(x), sparse_backend, x)
    ) == Vector{ForwardDiff.Dual{ForwardDiff.Tag{typeof(copyto!),Float64},Float64,1}}
end;