Add Speed Tests

.github/workflows/CI.yml

@@ -25,6 +25,7 @@ jobs:
           - SmokeXYOut
           - SmokeXYIn
           - Regression
+          - Speed
         version:
           - release
           - lts

test/Project.toml

@@ -6,6 +6,7 @@ DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
@@ -28,6 +29,7 @@ DataFrames = "1"
 DifferentiationInterface = "0.7"
 Distances = "0.10"
 Distributions = "0.25"
+Enzyme = "0.13"
 ExplicitImports = "1"
 ForwardDiff = "1"
 JET = "0.9, 0.10"

test/runtests.jl

@@ -5,6 +5,7 @@ import ADTypes,
     DifferentiationInterface,
     Distances,
     Distributions,
+    Enzyme,
     ExplicitImports,
     ForwardDiff,
     JET,
@@ -38,6 +39,10 @@ Test.@testset "Overall" begin
         include("regression_tests.jl")
     end
 
+    if GROUP == "All" || GROUP == "Speed"
+        include("speed_tests.jl")
+    end
+
     if GROUP == "All" || GROUP == "CheckByAqua"
         include("checkby_Aqua_tests.jl")
     end

test/smoke_tests.jl

@@ -21,49 +21,48 @@ Test.@testset "Smoke Tests" begin
     n_epochs_ = Int[2]
     data_types = Type{<:AbstractFloat}[Float32]
     devices = MLDataDevices.AbstractDevice[MLDataDevices.cpu_device()]
-    adtypes = ADTypes.AbstractADType[
-        ADTypes.AutoZygote(),
-        ADTypes.AutoForwardDiff(),
-        # ADTypes.AutoEnzyme(;
-        #     mode = Enzyme.set_runtime_activity(Enzyme.Reverse),
-        #     function_annotation = Enzyme.Const,
-        # ),
-        # ADTypes.AutoEnzyme(;
-        #     mode = Enzyme.set_runtime_activity(Enzyme.Forward),
-        #     function_annotation = Enzyme.Const,
-        # ),
+    adtypes = ADTypes.AbstractADType[ADTypes.AutoZygote(),
+    # ADTypes.AutoEnzyme(;
+    #     mode = Enzyme.set_runtime_activity(Enzyme.Reverse),
+    #     function_annotation = Enzyme.Const,
+    # ),
+    # ADTypes.AutoEnzyme(;
+    #     mode = Enzyme.set_runtime_activity(Enzyme.Forward),
+    #     function_annotation = Enzyme.Const,
+    # ),
+    # ADTypes.AutoForwardDiff(),
     ]
     compute_modes = ContinuousNormalizingFlows.ComputeMode[
         ContinuousNormalizingFlows.LuxVecJacMatrixMode(ADTypes.AutoZygote()),
-        ContinuousNormalizingFlows.LuxJacVecMatrixMode(ADTypes.AutoForwardDiff()),
         ContinuousNormalizingFlows.DIVecJacVectorMode(ADTypes.AutoZygote()),
-        ContinuousNormalizingFlows.DIJacVecVectorMode(ADTypes.AutoForwardDiff()),
         ContinuousNormalizingFlows.DIVecJacMatrixMode(ADTypes.AutoZygote()),
+        ContinuousNormalizingFlows.DIVecJacVectorMode(
+            ADTypes.AutoEnzyme(;
+                mode = Enzyme.set_runtime_activity(Enzyme.Reverse),
+                function_annotation = Enzyme.Const,
+            ),
+        ),
+        ContinuousNormalizingFlows.DIJacVecVectorMode(
+            ADTypes.AutoEnzyme(;
+                mode = Enzyme.set_runtime_activity(Enzyme.Forward),
+                function_annotation = Enzyme.Const,
+            ),
+        ),
+        ContinuousNormalizingFlows.DIVecJacMatrixMode(
+            ADTypes.AutoEnzyme(;
+                mode = Enzyme.set_runtime_activity(Enzyme.Reverse),
+                function_annotation = Enzyme.Const,
+            ),
+        ),
+        ContinuousNormalizingFlows.DIJacVecMatrixMode(
+            ADTypes.AutoEnzyme(;
+                mode = Enzyme.set_runtime_activity(Enzyme.Forward),
+                function_annotation = Enzyme.Const,
+            ),
+        ),
+        ContinuousNormalizingFlows.LuxJacVecMatrixMode(ADTypes.AutoForwardDiff()),
+        ContinuousNormalizingFlows.DIJacVecVectorMode(ADTypes.AutoForwardDiff()),
         ContinuousNormalizingFlows.DIJacVecMatrixMode(ADTypes.AutoForwardDiff()),
-        # ContinuousNormalizingFlows.DIVecJacVectorMode(
-        #     ADTypes.AutoEnzyme(;
-        #         mode = Enzyme.set_runtime_activity(Enzyme.Reverse),
-        #         function_annotation = Enzyme.Const,
-        #     ),
-        # ),
-        # ContinuousNormalizingFlows.DIJacVecVectorMode(
-        #     ADTypes.AutoEnzyme(;
-        #         mode = Enzyme.set_runtime_activity(Enzyme.Forward),
-        #         function_annotation = Enzyme.Const,
-        #     ),
-        # ),
-        # ContinuousNormalizingFlows.DIVecJacMatrixMode(
-        #     ADTypes.AutoEnzyme(;
-        #         mode = Enzyme.set_runtime_activity(Enzyme.Reverse),
-        #         function_annotation = Enzyme.Const,
-        #     ),
-        # ),
-        # ContinuousNormalizingFlows.DIJacVecMatrixMode(
-        #     ADTypes.AutoEnzyme(;
-        #         mode = Enzyme.set_runtime_activity(Enzyme.Forward),
-        #         function_annotation = Enzyme.Const,
-        #     ),
-        # ),
     ]
 
     Test.@testset "$device | $data_type | $compute_mode | ndata = $ndata | nvars = $nvars | inplace = $inplace | cond = $cond | planar = $planar | $omode | $mt" for device in

test/speed_tests.jl

@@ -0,0 +1,75 @@
+Test.@testset "Speed Tests" begin
+    compute_modes = ContinuousNormalizingFlows.ComputeMode[
+        ContinuousNormalizingFlows.LuxVecJacMatrixMode(ADTypes.AutoZygote()),
+        ContinuousNormalizingFlows.DIVecJacMatrixMode(ADTypes.AutoZygote()),
+        ContinuousNormalizingFlows.DIVecJacMatrixMode(
+            ADTypes.AutoEnzyme(;
+                mode = Enzyme.set_runtime_activity(Enzyme.Reverse),
+                function_annotation = Enzyme.Const,
+            ),
+        ),
+        ContinuousNormalizingFlows.DIJacVecMatrixMode(
+            ADTypes.AutoEnzyme(;
+                mode = Enzyme.set_runtime_activity(Enzyme.Forward),
+                function_annotation = Enzyme.Const,
+            ),
+        ),
+        ContinuousNormalizingFlows.DIJacVecMatrixMode(ADTypes.AutoForwardDiff()),
+        ContinuousNormalizingFlows.LuxJacVecMatrixMode(ADTypes.AutoForwardDiff()),
+    ]
+
+    Test.@testset "$compute_mode" for compute_mode in compute_modes
+        @show compute_mode
+
+        rng = StableRNGs.StableRNG(1)
+        nvars = 1
+        naugs = nvars
+        n_in = nvars + naugs
+        n = 2^10
+        nn = Lux.Chain(Lux.Dense(n_in => 3 * n_in, tanh), Lux.Dense(3 * n_in => n_in, tanh))
+
+        icnf = ContinuousNormalizingFlows.construct(
+            ContinuousNormalizingFlows.ICNF,
+            nn,
+            nvars,
+            naugs;
+            compute_mode,
+            tspan = (0.0f0, 1.0f0),
+            steer_rate = 1.0f-1,
+            λ₁ = 1.0f-2,
+            λ₂ = 1.0f-2,
+            λ₃ = 1.0f-2,
+            rng,
+            sol_kwargs = (;
+                save_everystep = false,
+                alg = OrdinaryDiffEqDefault.DefaultODEAlgorithm(),
+                sensealg = SciMLSensitivity.InterpolatingAdjoint(),
+            ),
+        )
+
+        data_dist = Distributions.Beta{Float32}(2.0f0, 4.0f0)
+        r = rand(icnf.rng, data_dist, nvars, n)
+        r = convert.(Float32, r)
+
+        df = DataFrames.DataFrame(transpose(r), :auto)
+        model = ContinuousNormalizingFlows.ICNFModel(icnf; batch_size = 0, n_epochs = 5)
+
+        mach = MLJBase.machine(model, df)
+        Test.@test !isnothing(MLJBase.fit!(mach))
+
+        @show only(MLJBase.report(mach).stats).time
+
+        d = ContinuousNormalizingFlows.ICNFDist(mach, ContinuousNormalizingFlows.TestMode())
+        actual_pdf = Distributions.pdf.(data_dist, r)
+        estimated_pdf = Distributions.pdf(d, r)
+
+        mad_ = Distances.meanad(estimated_pdf, actual_pdf)
+        msd_ = Distances.msd(estimated_pdf, actual_pdf)
+        tv_dis = Distances.totalvariation(estimated_pdf, actual_pdf) / n
+
+        @show mad_
+        @show msd_
+        @show tv_dis
+        Test.@test true
+    end
+end