feat: generalize indexing to all wrappers (#146)

* feat: generalize indexing to all wrappers

* test: use `PermuteDimsArray` to test parentindices

Author: avik-pal

src/TracedRArray.jl

@@ -34,10 +34,8 @@ ancestor(x::TracedRArray) = x
 ancestor(x::WrappedTracedRArray) = ancestor(parent(x))
 
 get_ancestor_indices(::TracedRArray, indices...) = indices
-function get_ancestor_indices(
-    x::SubArray{T,N,<:AnyTracedRArray{T,N}}, indices...
-) where {T,N}
-    return get_ancestor_indices(parent(x), Base.reindex(x.indices, indices)...)
+function get_ancestor_indices(x::WrappedTracedRArray, indices...)
+    return get_ancestor_indices(parent(x), Base.reindex(parentindices(x), indices)...)
 end
 
 Base.getindex(a::AnyTracedRScalar{T}) where {T} = a

test/wrapped_arrays.jl

@@ -113,3 +113,17 @@ end
 
     @test btranspose_badjoint_compiled(x_ra) ≈ btranspose_badjoint(x)
 end
+
+function bypass_permutedims(x)
+    x = PermutedDimsArray(x, (2, 1, 3))  # Don't use permutedims here
+    return view(x, 2:3, 1:2, :)
+end
+
+@testset "PermutedDimsArray" begin
+    x = rand(4, 4, 3)
+    x_ra = Reactant.to_rarray(x)
+
+    bypass_permutedims_compiled = @compile bypass_permutedims(x_ra)
+
+    @test bypass_permutedims_compiled(x_ra) ≈ bypass_permutedims(x)
+end