DifferentiationInterface.jl/DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/differentiate_with.jl at c765d67bfe29ce91c80ea64795259e3755d4de82 · 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
const NumberOrArray = Union{Number, AbstractArray{<:Number}}

# Mark DifferentiateWith with a range of context arities as primitives.
# For C contexts, the corresponding call tuple type is
# Tuple{DI.DifferentiateWith{C}, Any, Vararg{Any, C}}:
# one slot for the primal input x and C slots for contexts.
for C in 0:16
    @eval @is_primitive MinimalCtx Tuple{DI.DifferentiateWith{$C}, Vararg{Any, $(C + 1)}}
end
struct MooncakeDifferentiateWithError <: Exception
    F::Type
    X::Type
    Y::Type
    function MooncakeDifferentiateWithError(::F, ::X, ::Y) where {F, X, Y}
        return new(F, X, Y)
    end
end

function Base.showerror(io::IO, e::MooncakeDifferentiateWithError)
    return print(
        io,
        "MooncakeDifferentiateWithError: For the function type `$(e.F)` and input types `$(e.X)`, the output type `$(e.Y)` is currently not supported.",
    )
end

function Mooncake.rrule!!(
        dw::CoDual{<:DI.DifferentiateWith{C}},
        x::CoDual{<:Number},
        contexts::Vararg{CoDual, C}
    ) where {C}
    @assert tangent_type(typeof(dw)) == NoTangent
    primal_func = primal(dw)
    primal_x = primal(x)
    primal_contexts = map(primal, contexts)
    (; f, backend, context_wrappers) = primal_func
    y = zero_fcodual(f(primal_x, primal_contexts...))
    wrapped_primal_contexts = map(DI.call, context_wrappers, primal_contexts)

    # output is a vector, so we need to use the vector pullback
    function pullback_array!!(dy::NoRData)
        dx = DI.pullback(f, backend, primal_x, (y.dx,), wrapped_primal_contexts...) |> only
        @assert rdata(dx) isa rdata_type(tangent_type(typeof(primal_x)))
        rc = nanify_fdata_and_rdata!!(contexts...)
        return (NoRData(), rdata(dx), rc...)
    end

    # output is a scalar, so we can use the scalar pullback
    function pullback_scalar!!(dy::Number)
        dx = DI.pullback(f, backend, primal_x, (dy,), wrapped_primal_contexts...) |> only
        @assert rdata(dx) isa rdata_type(tangent_type(typeof(primal_x)))
        rc = nanify_fdata_and_rdata!!(contexts...)
        return (NoRData(), rdata(dx), rc...)
    end

    pullback = if primal(y) isa Number
        pullback_scalar!!
    elseif primal(y) isa AbstractArray
        pullback_array!!
    else
        throw(MooncakeDifferentiateWithError(primal_func, (primal_x, primal_contexts...), primal(y)))
    end

    return y, pullback
end

function Mooncake.rrule!!(
        dw::CoDual{<:DI.DifferentiateWith{C}},
        x::CoDual{<:AbstractArray{<:Number}},
        contexts::Vararg{CoDual, C}
    ) where {C}
    @assert tangent_type(typeof(dw)) == NoTangent
    primal_func = primal(dw)
    primal_x = primal(x)
    primal_contexts = map(primal, contexts)
    (; f, backend, context_wrappers) = primal_func
    y = zero_fcodual(f(primal_x, primal_contexts...))
    wrapped_primal_contexts = map(DI.call, context_wrappers, primal_contexts)

    # output is a vector, so we need to use the vector pullback
    function pullback_array!!(dy::NoRData)
        dx = DI.pullback(f, backend, primal_x, (y.dx,), wrapped_primal_contexts...) |> only
        @assert rdata(dx) isa rdata_type(tangent_type(typeof(primal_x)))
        x.dx .+= dx
        rc = nanify_fdata_and_rdata!!(contexts...)
        return (NoRData(), dy, rc...)
    end

    # output is a scalar, so we can use the scalar pullback
    function pullback_scalar!!(dy::Number)
        dx = DI.pullback(f, backend, primal_x, (dy,), wrapped_primal_contexts...) |> only
        @assert rdata(dx) isa rdata_type(tangent_type(typeof(primal_x)))
        x.dx .+= dx
        rc = nanify_fdata_and_rdata!!(contexts...)
        return (NoRData(), NoRData(), rc...)
    end

    pullback = if primal(y) isa Number
        pullback_scalar!!
    elseif primal(y) isa AbstractArray
        pullback_array!!
    else
        throw(MooncakeDifferentiateWithError(primal_func, (primal_x, primal_contexts...), primal(y)))
    end

    return y, pullback
end