Merge 5dbc35b into c768c0b

Author: prbzrg

src/utils.jl

@@ -6,7 +6,12 @@ function jacobian_batched(
     y = f(xs)
     z = similar(xs)
     ChainRulesCore.@ignore_derivatives fill!(z, zero(T))
-    res = Zygote.Buffer(xs, size(xs, 1), size(xs, 1), size(xs, 2))
+    res = Zygote.Buffer(
+        convert.(promote_type(eltype(xs), eltype(f.ps)), xs),
+        size(xs, 1),
+        size(xs, 1),
+        size(xs, 2),
+    )
     for i in axes(xs, 1)
         ChainRulesCore.@ignore_derivatives z[i, :] .= one(T)
         res[i, :, :] =
@@ -24,7 +29,12 @@ function jacobian_batched(
     y = f(xs)
     z = similar(xs)
     ChainRulesCore.@ignore_derivatives fill!(z, zero(T))
-    res = Zygote.Buffer(xs, size(xs, 1), size(xs, 1), size(xs, 2))
+    res = Zygote.Buffer(
+        convert.(promote_type(eltype(xs), eltype(f.ps)), xs),
+        size(xs, 1),
+        size(xs, 1),
+        size(xs, 2),
+    )
     for i in axes(xs, 1)
         ChainRulesCore.@ignore_derivatives z[i, :] .= one(T)
         res[:, i, :] = only(

test/smoke_tests.jl

@@ -36,6 +36,9 @@ Test.@testset "Smoke Tests" begin
         ContinuousNormalizingFlows.LuxVecJacMatrixMode(ADTypes.AutoZygote()),
         ContinuousNormalizingFlows.DIVecJacVectorMode(ADTypes.AutoZygote()),
         ContinuousNormalizingFlows.DIVecJacMatrixMode(ADTypes.AutoZygote()),
+        ContinuousNormalizingFlows.LuxJacVecMatrixMode(ADTypes.AutoForwardDiff()),
+        ContinuousNormalizingFlows.DIJacVecVectorMode(ADTypes.AutoForwardDiff()),
+        ContinuousNormalizingFlows.DIJacVecMatrixMode(ADTypes.AutoForwardDiff()),
         ContinuousNormalizingFlows.DIVecJacVectorMode(
             ADTypes.AutoEnzyme(;
                 mode = Enzyme.set_runtime_activity(Enzyme.Reverse),
@@ -60,9 +63,6 @@ Test.@testset "Smoke Tests" begin
                 function_annotation = Enzyme.Const,
             ),
         ),
-        ContinuousNormalizingFlows.LuxJacVecMatrixMode(ADTypes.AutoForwardDiff()),
-        ContinuousNormalizingFlows.DIJacVecVectorMode(ADTypes.AutoForwardDiff()),
-        ContinuousNormalizingFlows.DIJacVecMatrixMode(ADTypes.AutoForwardDiff()),
     ]
 
     Test.@testset "$device | $data_type | $compute_mode | ndata = $ndata | nvars = $nvars | inplace = $inplace | cond = $cond | planar = $planar | $omode | $mt" for device in
@@ -194,8 +194,12 @@ Test.@testset "Smoke Tests" begin
         Test.@test !isnothing(rand(d, ndata))
 
         Test.@testset "$adtype on loss" for adtype in adtypes
-            Test.@test !isnothing(DifferentiationInterface.gradient(diff_loss, adtype, ps))
-            Test.@test !isnothing(DifferentiationInterface.gradient(diff2_loss, adtype, r))
+            Test.@test !isnothing(DifferentiationInterface.gradient(diff_loss, adtype, ps)) broken =
+                GROUP != "All" &&
+                compute_mode.adback isa ADTypes.AutoEnzyme{<:Enzyme.ForwardMode}
+            Test.@test !isnothing(DifferentiationInterface.gradient(diff2_loss, adtype, r)) broken =
+                GROUP != "All" &&
+                compute_mode.adback isa ADTypes.AutoEnzyme{<:Enzyme.ForwardMode}
 
             Test.@testset "$n_epochs for fit" for n_epochs in n_epochs_
                 if cond

test/speed_tests.jl

@@ -2,6 +2,8 @@ Test.@testset "Speed Tests" begin
     compute_modes = ContinuousNormalizingFlows.ComputeMode[
         ContinuousNormalizingFlows.LuxVecJacMatrixMode(ADTypes.AutoZygote()),
         ContinuousNormalizingFlows.DIVecJacMatrixMode(ADTypes.AutoZygote()),
+        ContinuousNormalizingFlows.LuxJacVecMatrixMode(ADTypes.AutoForwardDiff()),
+        ContinuousNormalizingFlows.DIJacVecMatrixMode(ADTypes.AutoForwardDiff()),
         ContinuousNormalizingFlows.DIVecJacMatrixMode(
             ADTypes.AutoEnzyme(;
                 mode = Enzyme.set_runtime_activity(Enzyme.Reverse),
@@ -14,8 +16,6 @@ Test.@testset "Speed Tests" begin
                 function_annotation = Enzyme.Const,
             ),
         ),
-        ContinuousNormalizingFlows.DIJacVecMatrixMode(ADTypes.AutoForwardDiff()),
-        ContinuousNormalizingFlows.LuxJacVecMatrixMode(ADTypes.AutoForwardDiff()),
     ]
 
     Test.@testset "$compute_mode" for compute_mode in compute_modes
@@ -57,7 +57,9 @@ Test.@testset "Speed Tests" begin
         model = ContinuousNormalizingFlows.ICNFModel(icnf; batch_size = 0, n_epochs = 5)
 
         mach = MLJBase.machine(model, df)
-        Test.@test !isnothing(MLJBase.fit!(mach))
+        Test.@test !isnothing(MLJBase.fit!(mach)) broken =
+            GROUP != "All" &&
+            compute_mode.adback isa ADTypes.AutoEnzyme{<:Enzyme.ForwardMode}
 
         @show only(MLJBase.report(mach).stats).time