fix: handle MutableTangent and forward mode tangent leaks

Also convert leaked Mooncake.MutableTangent (e.g. MVector tangents)
and apply _maybe_to_primal in forward mode (pushforward) paths.

Author: sunxd3

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/forward_onearg.jl

@@ -41,7 +41,7 @@ function DI.value_and_pushforward(
             map(first_unwrap, contexts, prep.context_tangents)...,
         )
         y = first(y_and_dy)
-        dy = _copy_output(last(y_and_dy))
+        dy = _maybe_to_primal(last(y_and_dy), y)
         return y, dy
     end
     y = _copy_output(first(ys_and_ty[1]))

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/forward_twoarg.jl

@@ -55,7 +55,7 @@ function DI.value_and_pushforward(
             (x, dx),
             map(first_unwrap, contexts, prep.context_tangents)...,
         )
-        return _copy_output(new_dy)
+        return _maybe_to_primal(new_dy, y)
     end
     return y, ty
 end
@@ -93,7 +93,7 @@ function DI.value_and_pushforward!(
             (x, dx),
             map(first_unwrap, contexts, prep.context_tangents)...,
         )
-        copyto!(dy, new_dy)
+        copyto!(dy, _maybe_to_primal(new_dy, y))
     end
     return y, ty
 end

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/utils.jl

@@ -25,6 +25,7 @@ end
 # convert it to a primal-shaped value. No-op for already-converted results.
 _maybe_to_primal(tx, x) = _copy_output(tx)
 _maybe_to_primal(tx::Mooncake.Tangent, x) = tangent_to_user_primal(tx, x)
+_maybe_to_primal(tx::Mooncake.MutableTangent, x) = tangent_to_user_primal(tx, x)
 
 @inline maybe_getfield(mod, name::Symbol) =
     isdefined(mod, name) ? getfield(mod, name) : nothing