Fix reshape copy (#1323)

* Fix reshape copy

* Apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: wsmoses

src/TracedRArray.jl

@@ -747,6 +747,10 @@ function Base.copyto!(dest::TracedRArray{T,N}, src::TracedRArray{T2,N}) where {T
     return copyto!(dest, Ops.convert(TracedRArray{T,N}, src))
 end
 
+function Base.copyto!(dest::AnyTracedRArray, src::AnyTracedRArray)
+    return copyto!(dest, materialize_traced_array(src))
+end
+
 function Base.copyto!(dest::TracedRArray{T,N}, src::Array{T2,N}) where {T,T2,N}
     return copyto!(dest, TracedUtils.promote_to(TracedRArray{T2,N}, src))
 end

test/basic.jl

@@ -1061,6 +1061,23 @@ end
     @test Array(x_ra) == x
 end
 
+function reshapecopy!(x, y)
+    Base.copyto!(x, reshape(y, size(x)))
+    return nothing
+end
+@testset "copyto! Reshaped TracedRArray" begin
+    x = zeros(3, 4, 5)
+    y = collect(reshape(1:60, (3, 20)))
+
+    xr = Reactant.to_rarray(x)
+    yr = Reactant.to_rarray(y)
+
+    @jit reshapecopy!(xr, yr)
+
+    reshapecopy!(x, y)
+    @test Array(xr) == x
+end
+
 @testset "copy(::Broadcast.Broadcasted{ArrayStyle{ConcreteRArray}})" begin
     x_ra = Reactant.to_rarray(ones(4, 4))
     res = copy(Broadcast.broadcasted(-, Broadcast.broadcasted(+, x_ra, 1)))