DifferentiationInterface.jl/DifferentiationInterface/src/fallbacks/change_prep.jl at 55a3ddcd4badc09df64bbbb6e4c26571d9e888ab · 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
for op in [
    :derivative,
    :gradient,
    :jacobian,
    :second_derivative,
    :hessian,
    :pushforward,
    :pullback,
    :hvp,
]
    op! = Symbol(op, "!")
    val_and_op = if op == :second_derivative
        :value_derivative_and_second_derivative
    elseif op == :hessian
        :value_gradient_and_hessian
    elseif op == :hvp
        nothing
    else
        Symbol("value_and_", op)
    end
    val_and_op! = Symbol(val_and_op, "!")
    prep_op = Symbol("prepare_", op)
    prep_op_nokwarg = Symbol("prepare_", op, "_nokwarg")
    prep_op! = Symbol("prepare!_", op)
    prep_op_same_point = Symbol("prepare_", op, "_same_point")
    prep_op_same_point_nokwarg = Symbol("prepare_", op, "_same_point_nokwarg")
    P = if op == :derivative
        DerivativePrep
    elseif op == :gradient
        GradientPrep
    elseif op == :jacobian
        JacobianPrep
    elseif op == :second_derivative
        SecondDerivativePrep
    elseif op == :hessian
        HessianPrep
    elseif op == :pushforward
        PushforwardPrep
    elseif op == :pullback
        PullbackPrep
    elseif op == :hvp
        HVPPrep
    end

    if op in (:derivative, :gradient, :jacobian)
        # 1-arg
        @eval function $prep_op!(
            f::F, old_prep::$P, backend::AbstractADType, x, contexts::Vararg{Context,C}
        ) where {F,C}
            check_prep(f, old_prep, backend, x, contexts...)
            return $prep_op_nokwarg(is_strict(old_prep), f, backend, x, contexts...)
        end
        op == :gradient && continue
        # 2-arg
        @eval function $prep_op!(
            f!::F, y, old_prep::$P, backend::AbstractADType, x, contexts::Vararg{Context,C};
        ) where {F,C}
            check_prep(f!, y, old_prep, backend, x, contexts...)
            return $prep_op_nokwarg(is_strict(old_prep), f!, y, backend, x, contexts...)
        end

    elseif op in (:second_derivative, :hessian)
        # 1-arg
        @eval function $prep_op!(
            f::F, old_prep::$P, backend::AbstractADType, x, contexts::Vararg{Context,C};
        ) where {F,C}
            check_prep(f, old_prep, backend, x, contexts...)
            return $prep_op_nokwarg(is_strict(old_prep), f, backend, x, contexts...)
        end

    elseif op in (:pushforward, :pullback, :hvp)
        # 1-arg
        @eval function $prep_op!(
            f::F,
            old_prep::$P,
            backend::AbstractADType,
            x,
            seed::NTuple,
            contexts::Vararg{Context,C};
        ) where {F,C}
            check_prep(f, old_prep, backend, x, seed, contexts...)
            return $prep_op_nokwarg(is_strict(old_prep), f, backend, x, seed, contexts...)
        end
        @eval function $prep_op_same_point(
            f::F,
            prep::$P,
            backend::AbstractADType,
            x,
            seed::NTuple,
            contexts::Vararg{Context,C},
        ) where {F,C}
            check_prep(f, prep, backend, x, seed, contexts...)
            return prep
        end
        @eval function $prep_op_same_point_nokwarg(
            strict::Val,
            f::F,
            backend::AbstractADType,
            x,
            seed::NTuple,
            contexts::Vararg{Context,C};
        ) where {F,C}
            prep = $prep_op_nokwarg(strict, f, backend, x, seed, contexts...)
            return $prep_op_same_point(f, prep, backend, x, seed, contexts...)
        end
        op == :hvp && continue
        # 2-arg
        @eval function $prep_op!(
            f!::F,
            y,
            old_prep::$P,
            backend::AbstractADType,
            x,
            seed::NTuple,
            contexts::Vararg{Context,C},
        ) where {F,C}
            check_prep(f!, y, old_prep, backend, x, seed, contexts...)
            return $prep_op_nokwarg(
                is_strict(old_prep), f!, y, backend, x, seed, contexts...
            )
        end
        @eval function $prep_op_same_point(
            f!::F,
            y,
            prep::$P,
            backend::AbstractADType,
            x,
            seed::NTuple,
            contexts::Vararg{Context,C},
        ) where {F,C}
            check_prep(f!, y, prep, backend, x, seed, contexts...)
            return prep
        end
        @eval function $prep_op_same_point_nokwarg(
            strict::Val,
            f!::F,
            y,
            backend::AbstractADType,
            x,
            seed::NTuple,
            contexts::Vararg{Context,C};
        ) where {F,C}
            prep = $prep_op_nokwarg(strict, f!, y, backend, x, seed, contexts...)
            return $prep_op_same_point(f!, y, prep, backend, x, seed, contexts...)
        end
    end
end