fix: speed up Mooncake by avoiding tuple broadcasting

Author: gdalle

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/forward_onearg.jl

@@ -47,7 +47,7 @@ function DI.value_and_pushforward(
         return y, dy
     end
     y = first(ys_and_ty[1])
-    ty = last.(ys_and_ty)
+    ty = map(last, ys_and_ty)
     return y, ty
 end
 

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/onearg.jl

@@ -55,7 +55,7 @@ function DI.value_and_pullback(
         y, _copy_output(new_dx)
     end
     y = first(ys_and_tx[1])
-    tx = last.(ys_and_tx)
+    tx = map(last, ys_and_tx)
     return y, tx
 end