Merge pull request #116 from SciML/apindex

workaround Base changes for multi-indexing array partition

Author: ChrisRackauckas

Project.toml

@@ -1,7 +1,7 @@
 name = "RecursiveArrayTools"
 uuid = "731186ca-8d62-57ce-b412-fbd966d074cd"
 authors = ["Chris Rackauckas <[email protected]>"]
-version = "2.7.1"
+version = "2.7.2"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"

src/array_partition.jl

@@ -182,6 +182,15 @@ Base.@propagate_inbounds function Base.setindex!(A::ArrayPartition, v, i::Int)
   end
 end
 
+# workaround for https://github.com/SciML/RecursiveArrayTools.jl/issues/49
+function Base._unsafe_getindex(::IndexStyle, A::ArrayPartition, I::Vararg{Union{Real, AbstractArray}, N}) where N
+    # This is specifically not inlined to prevent excessive allocations in type unstable code
+    shape = Base.index_shape(I...)
+    dest = similar(A.x[1], shape)
+    Base._unsafe_getindex!(dest, A, I...) # usually a generated function, don't allow it to impact inference result
+    return dest
+end
+
 """
     setindex!(A::ArrayPartition, v, i::Int, j...)
 

test/basic_indexing.jl

@@ -119,3 +119,8 @@ ref = sqrt.(abs.(x .* x))
 fill!(mulX, 0)
 mulX .= sqrt.(abs.(testva .* testvb))
 @test mulX == ref
+
+# https://github.com/SciML/RecursiveArrayTools.jl/issues/49
+a = ArrayPartition(1:5, 1:6)
+a[1:8]
+a[[1,3,8]]