Test DI, support Mooncake

Author: timholy

Project.toml

@@ -8,18 +8,24 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 
 [weakdeps]
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 
 [extensions]
 ThickNumbersForwardDiffExt = "ForwardDiff"
+ThickNumbersMooncakeExt = "Mooncake"
 
 [compat]
+DifferentiationInterface = "0.6"
 ForwardDiff = "0.10, 1"
 LinearAlgebra = "1"
+Mooncake = "0.4"
 julia = "1.9"
 
 [extras]
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["ForwardDiff", "Test"]
+test = ["DifferentiationInterface", "ForwardDiff", "Mooncake", "Test"]

ThickNumbersInterfaceTests/test/testpackages/IntervalArith/src/IntervalArith.jl

@@ -15,6 +15,7 @@ end
 Interval(lo, hi) = Interval(promote(lo, hi)...)
 Interval{T}(iv::Interval) where T = Interval{T}(iv.lo, iv.hi)
 Interval{T}(x::Number) where T = Interval{T}(x, x)
+Interval{T}(nt::@NamedTuple{lo::T, hi::T}) where T = Interval{T}(nt.lo, nt.hi)   # needed by Mooncake
 
 ThickNumbers.loval(x::Interval) = x.lo
 ThickNumbers.hival(x::Interval) = x.hi
@@ -44,4 +45,6 @@ end
 Base.abs2(x::Interval) = Interval(mig(x)^2, mag(x)^2)
 Base.sqrt(x::Interval) = Interval(sqrt(loval(x)), sqrt(hival(x)))
 
+Base.conj(x::Interval{T}) where T = Interval{T}(conj(x.lo), conj(x.hi))   # needed for Enzyme
+
 end # module IntervalArith

ext/ThickNumbersMooncakeExt.jl

@@ -0,0 +1,12 @@
+module ThickNumbersMooncakeExt
+
+using ThickNumbers
+using Mooncake
+
+Mooncake.tangent_type(::Type{TN}) where TN<:ThickNumber = TN
+Mooncake.fdata_type(::Type{TN}) where TN<:ThickNumber = Mooncake.NoFData
+Mooncake.rdata_type(::Type{TN}) where TN<:ThickNumber = TN
+Mooncake.zero_rdata(x::ThickNumber) = zero(x)
+Mooncake.increment_internal!!(::Mooncake.IncCache, x::TN, y::TN) where TN<:ThickNumber = x + y
+
+end

test/extensions/di.jl

@@ -0,0 +1,41 @@
+using ThickNumbers
+using DifferentiationInterface
+using ForwardDiff
+# using Enzyme: EnzymeCore
+using Mooncake
+
+using Test
+
+include(joinpath(dirname(@__DIR__), "setpath.jl"))
+
+using IntervalArith
+
+@testset "DifferentiationInterface" begin
+    @test isempty(detect_ambiguities(ThickNumbers))
+    @test isempty(detect_ambiguities(IntervalArith))
+
+    for backend in (AutoForwardDiff(), #=AutoEnzyme(mode=EnzymeCore.Forward), AutoEnzyme(mode=EnzymeCore.Reverse),=# AutoMooncake(config=nothing))
+        a, b = Interval(1, 2), Interval(0, 0.1)
+        f1(t) = a + t*b
+        f2(x) = a + abs2(x)/2
+
+        df1(t) = derivative(f1, backend, t)
+        df2(x) = derivative(f2, backend, x)
+        @test df1(0.5) ≐ b
+        @test df2(b) ⩪ b
+        ddf2(x) = derivative(df2, backend, x)
+        @test ddf2(b) ≐ 1
+
+        # abs
+        dabs(x) = derivative(abs, backend, x)
+        ddabs(x) = derivative(dabs, backend, x)
+        dddabs(x) = derivative(ddabs, backend, x)
+        @test dabs(Interval(1.0, 2.0)) === Interval(1.0, 1.0)
+        @test ddabs(Interval(1.0, 2.0)) === Interval(0.0, 0.0)
+        @test dddabs(Interval(1.0, 2.0)) === Interval(0.0, 0.0)
+        @test dabs(Interval(-1.0, 2.0)) === Interval(-1.0, 1.0)
+        @test ddabs(Interval(-1.0, 2.0)) === Interval(0.0, Inf)
+        abs3 = dddabs(Interval(-1.0, 2.0))
+        @test abs3 === Interval(-Inf, Inf) || isnan_tn(abs3)
+    end
+end

test/extensions/runtests.jl

@@ -1 +1,2 @@
 include("forwarddiff.jl")
+include("di.jl")