DifferentiationInterface.jl/DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/DifferentiationInterfaceMooncakeExt.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
module DifferentiationInterfaceMooncakeExt

using ADTypes: ADTypes, AutoMooncake, AutoMooncakeForward
import DifferentiationInterface as DI
using Mooncake:
    Mooncake,
    CoDual,
    Config,
    Dual,
    prepare_derivative_cache,
    prepare_gradient_cache,
    prepare_pullback_cache,
    primal,
    tangent,
    tangent_type,
    value_and_derivative!!,
    value_and_gradient!!,
    value_and_pullback!!,
    zero_dual,
    zero_tangent,
    zero_rdata,
    rdata_type,
    fdata,
    rdata,
    tangent_type,
    NoTangent,
    @is_primitive,
    zero_fcodual,
    MinimalCtx,
    NoFData,
    NoRData,
    primal,
    _copy_output,
    _copy_to_output!!,
    tangent_to_primal!!,
    increment!!

const AnyAutoMooncake{C} = Union{AutoMooncake{C}, AutoMooncakeForward{C}}

DI.check_available(::AnyAutoMooncake{C}) where {C} = true
DI.inner_preparation_behavior(::AutoMooncakeForward) = DI.PrepareInnerSimple()

include("utils.jl")
include("onearg.jl")
include("twoarg.jl")
include("forward_onearg.jl")
include("forward_twoarg.jl")
include("differentiate_with.jl")

end