update tests to new interface

Author: zuhengxu

test/Project.toml

@@ -1,11 +1,11 @@
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
-DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
-Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"

test/ad.jl

@@ -1,21 +1,22 @@
-@testset "AD correctness" begin
-    f(x) = sum(abs2, x)
+@testset "DI.AD with context wrapper" begin
+    f(x, y, z) = sum(abs2, x .+ y .+ z)
 
     @testset "$T" for T in [Float32, Float64]
         x = randn(T, 10)
+        y = randn(T, 10)
+        z = randn(T, 10)
         chunksize = size(x, 1)
 
         @testset "$at" for at in [
             ADTypes.AutoZygote(),
             ADTypes.AutoForwardDiff(; chunksize=chunksize),
             ADTypes.AutoForwardDiff(),
             ADTypes.AutoReverseDiff(false),
-            ADTypes.AutoEnzyme(),
+            ADTypes.AutoMooncake(; config=ADTypes.Mooncake.Config()),
         ]
-            out = DiffResults.GradientResult(x)
-            NormalizingFlows.value_and_gradient!(at, f, x, out)
-            @test DiffResults.value(out) ≈ f(x)
-            @test DiffResults.gradient(out) ≈ 2x
+            value, grad = NormalizingFlows._value_and_gradient(f, at, x, y, z)
+            @test DiffResults.value(out) ≈ f(x, y, z)
+            @test DiffResults.gradient(out) ≈ 2 * (x .+ y .+ z)
         end
     end
 end
@@ -25,7 +26,7 @@ end
         ADTypes.AutoZygote(),
         ADTypes.AutoForwardDiff(),
         ADTypes.AutoReverseDiff(false),
-        # ADTypes.AutoEnzyme(), # not working now
+        ADTypes.AutoMooncake(; config=ADTypes.Mooncake.Config()),
     ]
         @testset "$T" for T in [Float32, Float64]
             μ = 10 * ones(T, 2)
@@ -38,15 +39,15 @@ end
 
             sample_per_iter = 10
             θ, re = Optimisers.destructure(flow)
-            out = DiffResults.GradientResult(θ)
 
             # check grad computation for elbo
-            NormalizingFlows.grad!(
-                Random.default_rng(), at, elbo, θ, re, out, logp, sample_per_iter
+            loss(θ, args...) = -NormalizingFlows.elbo(re(θ), args...)
+            value, grad = NormalizingFlows._value_and_gradient(
+                loss, at, θ, logp, randn(T, 2, sample_per_iter)
             )
 
-            @test DiffResults.value(out) != nothing
-            @test all(DiffResults.gradient(out) .!= nothing)
+            @test !isnothing(value)
+            @test all(grad .!= nothing)
         end
     end
-end
+end

test/interface.jl

@@ -4,8 +4,8 @@
         ADTypes.AutoZygote(),
         ADTypes.AutoForwardDiff(; chunksize=chunksize),
         ADTypes.AutoForwardDiff(),
-        ADTypes.AutoReverseDiff(false),
-        # ADTypes.AutoEnzyme(), # doesn't work for Enzyme
+        ADTypes.AutoReverseDiff(),
+        ADTypes.AutoMooncake(; config = ADTypes.Mooncake.Config()),
     ]
         @testset "$T" for T in [Float32, Float64]
             μ = 10 * ones(T, 2)
@@ -44,4 +44,4 @@
             @test el_trained > -1
         end
     end
-end
+end

test/objectives.jl

@@ -9,9 +9,10 @@
         flow = Bijectors.transformed(q₀, Bijectors.Shift(μ) ∘ Bijectors.Scale(sqrt.(Σ)))
 
         x = randn(T, 2)
+        rng = Random.default_rng()
 
         @testset "elbo" begin
-            el = elbo(Random.default_rng(), flow, logp, 10)
+            el = elbo(rng, flow, logp, 10)
 
             @test abs(el) ≤ 1e-5
             @test logpdf(flow, x) + el ≈ logp(x)
@@ -20,8 +21,8 @@
         @testset "likelihood" begin
             sample_trained = rand(flow, 1000)
             sample_untrained = rand(q₀, 1000)
-            llh_trained = NormalizingFlows.loglikelihood(flow, sample_trained)
-            llh_untrained = NormalizingFlows.loglikelihood(flow, sample_untrained)
+            llh_trained = NormalizingFlows.loglikelihood(rng, flow, sample_trained)
+            llh_untrained = NormalizingFlows.loglikelihood(rng, flow, sample_untrained)
 
             @test llh_trained > llh_untrained
         end

test/runtests.jl

@@ -3,10 +3,11 @@ using Distributions
 using Bijectors, Optimisers
 using LinearAlgebra
 using Random
-using ADTypes, DiffResults
+using ADTypes
+import DifferentiationInterface as DI
 using ForwardDiff, Zygote, Enzyme, ReverseDiff
 using Test
 
 include("ad.jl")
 include("objectives.jl")
-include("interface.jl")
+include("interface.jl")