fix: define getindexing into sub reshaped array

EnzymeAD · Jan 17, 2025 · 46029f4 · 46029f4
1 parent 532fd73
commit 46029f4
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/src/TracedRArray.jl b/src/TracedRArray.jl
@@ -112,7 +112,9 @@ end
 
 function Base.getindex(a::TracedRArray{T,N}, indices) where {T,N}
     if !(indices isa TracedRArray)
-        indices = TracedUtils.promote_to(TracedRArray{Int,1}, collect(indices))
+        indices = collect(indices)
+        eltype(indices) <: CartesianIndex && (indices = LinearIndices(size(a))[indices])
+        indices = TracedUtils.promote_to(TracedRArray{Int,1}, indices)
     end
     return Ops.gather_getindex(a, scalar_index_to_cartesian(indices, size(a)))
 end

diff --git a/src/TracedUtils.jl b/src/TracedUtils.jl
@@ -60,6 +60,15 @@ function set_mlir_data!(
     return x
 end
 
+function get_ancestor_indices(
+    x::WrappedReshapedArray{TracedRNumber{T},N,TracedRArray{T,M}}, indices...
+) where {T,N,M}
+    cartesian_indices = CartesianIndex.(indices...)
+    linear_indices = LinearIndices(size(x))[cartesian_indices]
+    parent_cartesian_indices = CartesianIndices(size(parent(x)))[linear_indices]
+    return (parent_cartesian_indices,)
+end
+
 function set_mlir_data!(
     x::PermutedDimsArray{TracedRNumber{T},N,perm,iperm,TracedRArray{T,N}}, data
 ) where {T,N,perm,iperm}