Fix permutedims! to support AnyTracedRArray destinations (including ReshapedArray)

Agent-Logs-Url: https://github.com/EnzymeAD/Reactant.jl/sessions/36ff73e5-274e-4c26-bbf9-acf3130b3cb0

Co-authored-by: avik-pal <30564094+avik-pal@users.noreply.github.com>

Author: Copilot

src/TracedRArray.jl

@@ -1338,7 +1338,7 @@ function Base.permutedims(A::AnyTracedRArray{T,N}, perm) where {T,N}
     return @opcall transpose(materialize_traced_array(A), Int64[perm...])
 end
 
-function Base.permutedims!(dest::TracedRArray, src::AnyTracedRArray, perm)
+function Base.permutedims!(dest::AnyTracedRArray, src::AnyTracedRArray, perm)
     result = @opcall transpose(materialize_traced_array(src), Int64[perm...])
     TracedUtils.set_mlir_data!(dest, result.mlir_data)
     return dest

test/core/wrapped_arrays.jl

@@ -294,3 +294,19 @@ end
     @test @jit(view_transpose(x_ra)) ≈ view_transpose(x)
     @test @jit(view_diagonal(x_ra)) ≈ view_diagonal(x)
 end
+
+function permutedims!_reshaped(A, B)
+    A_reshaped = reshape(A, 2, 2, 2)
+    B_reshaped = reshape(B, 2, 2, 2)
+    permutedims!(A_reshaped, B_reshaped, (2, 3, 1))
+    return A
+end
+
+@testset "permutedims! on reshaped arrays" begin
+    A = randn(Float32, 8)
+    B = randn(Float32, 8)
+    A_ra = Reactant.to_rarray(A)
+    B_ra = Reactant.to_rarray(B)
+
+    @test Array(@jit(permutedims!_reshaped(A_ra, B_ra))) ≈ permutedims!_reshaped(A, B)
+end