Convert to tuple if necessary

JuliaDiff · Jan 10, 2022 · bc3c321 · bc3c321
1 parent 8b513d0
commit bc3c321
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/test/defaults.jl b/test/defaults.jl
@@ -24,7 +24,8 @@ FDMBackend2() = FDMBackend2(central_fdm(5, 1))
 const fdm_backend2 = FDMBackend2()
 AD.@primitive function pushforward_function(ab::FDMBackend2, f, xs...)
     return function (vs)
-        FDM.jvp(ab.alg, f, tuple.(xs, vs)...)
+        ws = FDM.jvp(ab.alg, f, tuple.(xs, vs)...)
+        return length(xs) == 1 ? (ws,) : ws
     end
 end