_copy_temp() to replace Base.copy() usage.

src/interface.jl

@@ -16,7 +16,7 @@
     out, pb!! = rule(fx_fwds...)
     @assert _typeof(tangent(out)) == fdata_type(T)
     increment!!(tangent(out), fdata(ȳ))
-    v = y_cache === nothing ? copy(primal(out)) : _copy!!(y_cache, primal(out))
+    v = y_cache === nothing ? _copy_temp(primal(out)) : _copy!!(y_cache, primal(out))
     return v, tuple_map((f, r) -> tangent(fdata(tangent(f)), r), fx, pb!!(rdata(ȳ)))
 end
 
@@ -238,6 +238,63 @@
 
 const _BuiltinArrays = @static VERSION >= v"1.11" ? Union{Array,Memory} : Array
 
+# explicit for svec
+function _copy_temp(x::P) where {P<:SimpleVector}
+    return Core.svec([map(_copy_temp, x_sub) for x_sub in x]...)
+end
+
+# Array, Memory
+function _copy_temp(x::P) where {P<:_BuiltinArrays}
+    temp = P(undef, size(x)...)
+    @inbounds for i in eachindex(temp)
+        isassigned(x, i) && (temp[i] = _copy_temp(x[i]))
+    end
+    return temp
+end
+
+# Tuple, NamedTuple
+function _copy_temp(x::P) where {P<:Union{Tuple,NamedTuple}}
+    return map(_copy_temp, x)
+end
+
+# mutable composite types, bitstype
+function _copy_temp(x::P) where {P}
+    isbitstype(P) && return x
+    nf = nfields(P)
+
+    if ismutable(x)
+        temp = ccall(:jl_new_struct_uninit, Any, (Any,), P)
+        for x_sub in 1:nf
+            if isdefined(x, x_sub)
+                ccall(
+                    :jl_set_nth_field,
+                    Cvoid,
+                    (Any, Csize_t, Any),
+                    temp,
+                    x_sub - 1,
+                    _copy_temp(getfield(x, x_sub)),
+                )
+            end
+        end
+
+        return temp
+    else
+        flds = Vector{Any}(undef, nf)
+        for x_sub in 1:nf
+            if isdefined(x, x_sub)
+                flds[x_sub] = _copy_temp(getfield(x, x_sub))
+            else
+                nf = x_sub - 1  # Assumes if a undefined field is found, all subsequent fields are undefined.
+                break
+            end
+        end
+
+        # when immutable struct object created by non initializing inner constructor. (Base.deepcopy misses this out)
+        !isassigned(flds, 1) && return x
+        return ccall(:jl_new_structv, Any, (Any, Ptr{Any}, UInt32), P, flds, nf)
+    end
+end
+
 function __exclude_unsupported_output_internal!(
     y::T, address_set::Set{UInt}
 ) where {T<:_BuiltinArrays}
@@ -291,7 +348,7 @@
     __exclude_unsupported_output(y)
 
     # Construct cache for output. Check that `copy!`ing appears to work.
-    y_cache = copy(primal(y))
+    y_cache = _copy_temp(primal(y))
     return Cache(rule, _copy!!(y_cache, primal(y)), tangents)
 end
 

test/interface.jl

@@ -175,5 +175,61 @@ end
                 @test isa(err, Mooncake.ValueAndPullbackReturnTypeError)
             end
         end
+
+        # Test for new copy function
+        @testset for i in eachindex(additional_test_set)
+            original = additional_test_set[i][2]
+            try
+                if isnothing(Mooncake.__exclude_unsupported_output(original))
+                    test_copy = Mooncake._copy_temp(original)
+
+                    function comparisons(
+                        original::P, test_copy::P
+                    ) where {P<:Mooncake._BuiltinArrays}
+                        @test original !== test_copy
+                        @test size(original) == size(test_copy)
+
+                        # Value caching for pure immutable Types!
+                        for i in eachindex(test_copy)
+                            if !isassigned(test_copy, i)
+                                @test !isassigned(original, i)
+                            else
+                                comparisons(original[i], test_copy[i])
+                            end
+                        end
+                    end
+
+                    function comparisons(original::P, test_copy::P) where {P}
+                        (isbitstype(P) && !isnothing(original) && isnan(original)) &&
+                            return @test isnan(test_copy)
+                        isbitstype(P) && return @test test_copy == original
+
+                        fields_copy = [
+                            if !isdefined(test_copy, name)
+                                nothing
+                            else
+                                getfield(test_copy, name)
+                            end for name in fieldnames(typeof(test_copy))
+                        ]
+                        fields_orig = [
+                            !isdefined(original, name) ? nothing : getfield(original, name)
+                            for name in fieldnames(typeof(original))
+                        ]
+
+                        return comparisons(fields_orig, fields_copy)
+                    end
+
+                    @test typeof(test_copy) == typeof(original)
+                    # isbitstypes with same values are stored in the same address (Value Caching).
+                    if isbitstype(typeof(original))
+                        @test test_copy == original
+                    else
+                        comparisons(original, test_copy)
+                    end
+                end
+            catch err
+                @test isa(err, Mooncake.ValueAndPullbackReturnTypeError)
+            end
+        end
     end
 end