Fix size-1 Jacobian in Enzyme

gdalle · Jul 22, 2024 · 25797a8 · 25797a8
1 parent 5419f9a
commit 25797a8
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 3 deletions.
diff --git a/...ationInterface/ext/DifferentiationInterfaceEnzymeExt/DifferentiationInterfaceEnzymeExt.jl b/...ationInterface/ext/DifferentiationInterfaceEnzymeExt/DifferentiationInterfaceEnzymeExt.jl
@@ -36,7 +36,8 @@ using Enzyme:
     gradient!,
     jacobian,
     make_zero,
-    make_zero!
+    make_zero!,
+    onehot
 
 struct AutoDeferredEnzyme{M} <: ADTypes.AbstractADType
     mode::M

diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_onearg.jl
@@ -102,7 +102,11 @@ end
 
 function DI.prepare_jacobian(f, backend::AutoEnzyme{<:Union{ForwardMode,Nothing}}, x)
     B = pick_batchsize(backend, length(x))
-    shadow = chunkedonehot(x, Val(B))
+    if B == 1
+        shadow = onehot(x)
+    else
+        shadow = chunkedonehot(x, Val(B))
+    end
     return EnzymeForwardOneArgJacobianExtras{B,typeof(shadow)}(shadow)
 end
 

diff --git a/DifferentiationInterfaceTest/src/scenarios/default.jl b/DifferentiationInterfaceTest/src/scenarios/default.jl
@@ -28,12 +28,40 @@ function num_to_num_scenarios_onearg(x::Number; dx::Number, dy::Number)
     der2 = num_to_num_second_derivative(x)
 
     # everyone out of place
-    return [
+    scens = Scenario[
         PushforwardScenario(f; x, y, dx, dy=dy_from_dx, nb_args, place),
         PullbackScenario(f; x, y, dy, dx=dx_from_dy, nb_args, place),
         DerivativeScenario(f; x, y, der, nb_args, place),
         SecondDerivativeScenario(f; x, y, der, der2, nb_args, place),
     ]
+
+    # add scenarios [x] -> [y] to test 1-sized everything
+    f_vec(x) = [f(only(x))]
+    f_vec!(y, x) = y[only(eachindex(y))] = f(only(x))
+
+    for place in (:outofplace, :inplace)
+        append!(
+            scens,
+            [
+                PushforwardScenario(
+                    f_vec; x=[x], y=[y], dx=[dx], dy=[dy_from_dx], nb_args=1, place
+                ),
+                PushforwardScenario(
+                    f_vec!; x=[x], y=[y], dx=[dx], dy=[dy_from_dx], nb_args=2, place
+                ),
+                PullbackScenario(
+                    f_vec; x=[x], y=[y], dy=[dy], dx=[dx_from_dy], nb_args=1, place
+                ),
+                PullbackScenario(
+                    f_vec!; x=[x], y=[y], dy=[dy], dx=[dx_from_dy], nb_args=2, place
+                ),
+                JacobianScenario(f_vec; x=[x], y=[y], jac=[der;;], nb_args=1, place),
+                JacobianScenario(f_vec!; x=[x], y=[y], jac=[der;;], nb_args=2, place),
+            ],
+        )
+    end
+
+    return scens
 end
 
 ## Number to array