that wasnt actually a bug in Enzyme, its just really confusing

JuliaDiff · Sep 19, 2024 · 81dc141 · 81dc141
1 parent d524893
commit 81dc141
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/reverse_onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/reverse_onearg.jl
@@ -220,7 +220,7 @@ function DI.jacobian(
     backend::AutoEnzyme{<:ReverseMode,Nothing},
     x,
 ) where {M,B}
-    J = jacobian(mode_noprimal(backend), f, x; n_outs=Val(M), chunk=Val(B))[1]
+    J = jacobian(mode_noprimal(backend), f, x; n_outs=Val((M,)), chunk=Val(B))[1]
     flatjac(x, J)
 end
 
@@ -230,8 +230,7 @@ function DI.value_and_jacobian(
     backend::AutoEnzyme{<:ReverseMode,Nothing},
     x,
 ) where {M,B}
-    #TODO: right now this is giving different outputs if you pass args Enzyme #1863
-    jac = jacobian(mode_withprimal(backend), f, x)#; n_outs=Val(M), chunk=Val(B))
+    jac = jacobian(mode_withprimal(backend), f, x; n_outs=Val((M,)), chunk=Val(B))
     return jac.val, flatjac(x, jac.derivs[1])
 end