Merge branch 'master' into ChrisRackauckas-patch-1

Tokazama · web-flow · commit d6ee3596ac7f · 2022-06-14T06:16:32.000-04:00
diff --git a/.github/workflows/Downstream.yml b/.github/workflows/Downstream.yml
@@ -22,7 +22,11 @@ jobs:
           - {user: SciML, repo: DiffEqBase.jl, group: Core}
           - {user: SciML, repo: DiffEqBase.jl, group: Downstream}
           - {user: SciML, repo: DiffEqBase.jl, group: Downstream2}
-          - {user: SciML, repo: DiffEqSensitivity.jl, group: Core}
+          - {user: SciML, repo: DiffEqSensitivity.jl, group: Core1}
+          - {user: SciML, repo: DiffEqSensitivity.jl, group: Core2}
+          - {user: SciML, repo: DiffEqSensitivity.jl, group: Core3}
+          - {user: SciML, repo: DiffEqSensitivity.jl, group: Core4}
+          - {user: SciML, repo: DiffEqSensitivity.jl, group: Core5}
           - {user: SciML, repo: OrdinaryDiffEq.jl, group: Core}
           - {user: SciML, repo: OrdinaryDiffEq.jl, group: Interface}
           - {user: SciML, repo: DelayDiffEq.jl, group: Interface}
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "ArrayInterface"
 uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
-version = "6.0.14"
+version = "6.0.16"
 
 [deps]
 ArrayInterfaceCore = "30b0a656-2188-435a-8636-2ec0e6a096e2"
diff --git a/lib/ArrayInterfaceCore/Project.toml b/lib/ArrayInterfaceCore/Project.toml
@@ -1,6 +1,6 @@
 name = "ArrayInterfaceCore"
 uuid = "30b0a656-2188-435a-8636-2ec0e6a096e2"
-version = "0.1.10"
+version = "0.1.11"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
diff --git a/lib/ArrayInterfaceCore/src/ArrayInterfaceCore.jl b/lib/ArrayInterfaceCore/src/ArrayInterfaceCore.jl
@@ -548,6 +548,20 @@ ndims_index(@nospecialize T::Type{<:Base.LogicalIndex}) = ndims(fieldtype(T, :ma
 ndims_index(T::Type) = 1
 ndims_index(@nospecialize(i)) = ndims_index(typeof(i))
 
+"""
+    ndims_shape(::Type{I}) -> Union{Int,Tuple{Vararg{Int}}}
+
+Returns the number of dimension that are represented in shape of the returned array when
+indexing with an instance of `I`.
+"""
+ndims_shape(T::DataType) = ndims_index(T)
+ndims_shape(::Type{Colon}) = 1
+ndims_shape(T::Type{<:Base.AbstractCartesianIndex{N}}) where {N} = ntuple(zero, Val{N}())
+ndims_shape(@nospecialize T::Type{<:CartesianIndices}) = ntuple(one, Val{ndims(T)}())
+ndims_shape(@nospecialize T::Type{<:Number}) = 0
+ndims_shape(@nospecialize T::Type{<:AbstractArray}) = ndims(T)
+ndims_shape(x) = ndims_shape(typeof(x))
+
 """
     instances_do_not_alias(::Type{T}) -> Bool
 
diff --git a/lib/ArrayInterfaceCore/test/runtests.jl b/lib/ArrayInterfaceCore/test/runtests.jl
@@ -268,6 +268,14 @@ end
     @test @inferred(ArrayInterfaceCore.ndims_index(1)) == 1
 end
 
+@testset "ndims_shape" begin
+    @test @inferred(ArrayInterfaceCore.ndims_shape(1)) === 0
+    @test @inferred(ArrayInterfaceCore.ndims_shape(:)) === 1
+    @test @inferred(ArrayInterfaceCore.ndims_shape(CartesianIndex(1, 2))) === (0, 0)
+    @test @inferred(ArrayInterfaceCore.ndims_shape(CartesianIndices((2,2)))) === (1, 1)
+    @test @inferred(ArrayInterfaceCore.ndims_shape([1 1])) === 2
+end
+
 @testset "indices_do_not_alias" begin
   @test ArrayInterfaceCore.instances_do_not_alias(Float64)
   @test !ArrayInterfaceCore.instances_do_not_alias(Matrix{Float64})
@@ -285,4 +293,3 @@ end
   @test !ArrayInterfaceCore.indices_do_not_alias(typeof(view(fill(rand(4,4),4,4)', 2:3, 1:2)))
   @test !ArrayInterfaceCore.indices_do_not_alias(typeof(view(rand(4,4)', StepRangeLen(1,0,5), 1:2)))
 end
-
diff --git a/lib/ArrayInterfaceOffsetArrays/Project.toml b/lib/ArrayInterfaceOffsetArrays/Project.toml
@@ -1,6 +1,6 @@
 name = "ArrayInterfaceOffsetArrays"
 uuid = "015c0d05-e682-4f19-8f0a-679ce4c54826"
-version = "0.1.2"
+version = "0.1.5"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
diff --git a/lib/ArrayInterfaceOffsetArrays/src/ArrayInterfaceOffsetArrays.jl b/lib/ArrayInterfaceOffsetArrays/src/ArrayInterfaceOffsetArrays.jl
@@ -27,6 +27,7 @@ end
 function ArrayInterface.axes_types(::Type{T}) where {T<:OffsetArrays.OffsetArray}
     Static.eachop_tuple(_offset_axis_type, Static.nstatic(Val(ndims(T))), ArrayInterface.parent_type(T))
 end
+ArrayInterface.strides(A::OffsetArray) = ArrayInterface.strides(parent(A))
 function ArrayInterface.known_offsets(::Type{A}) where {A<:OffsetArrays.OffsetArray}
     ntuple(identity -> nothing, Val(ndims(A)))
 end
@@ -44,5 +45,14 @@ end
     d = ArrayInterface.to_dims(A, dim)
     OffsetArrays.IdOffsetRange(ArrayInterface.axes(parent(A), d), relative_offsets(A, d))
 end
+function ArrayInterface.stride_rank(T::Type{<:OffsetArray})
+  ArrayInterface.stride_rank(ArrayInterface.parent_type(T))
+end
+function ArrayInterface.dense_dims(T::Type{<:OffsetArray})
+    ArrayInterface.dense_dims(ArrayInterface.parent_type(T))
+end
+function ArrayInterface.contiguous_axis(T::Type{<:OffsetArray})
+  ArrayInterface.contiguous_axis(ArrayInterface.parent_type(T))
+end
 
 end # module
diff --git a/lib/ArrayInterfaceOffsetArrays/test/runtests.jl b/lib/ArrayInterfaceOffsetArrays/test/runtests.jl
@@ -11,6 +11,9 @@ Op = PermutedDimsArray(O,(3,1,2));
 @test @inferred(ArrayInterface.offsets(O)) === (4, 8, 11)
 @test @inferred(ArrayInterface.offsets(Op)) === (11, 4, 8)
 
+@test @inferred(ArrayInterface.to_indices(O, (:, :, :))) == (4:6, 8:11, 11:15)
+@test @inferred(ArrayInterface.to_indices(Op, (:, :, :))) == (11:15, 4:6, 8:11)
+
 @test @inferred(ArrayInterface.offsets((1,2,3))) === (StaticInt(1),)
 o = OffsetArray(vec(A), 8);
 @test @inferred(ArrayInterface.offset1(o)) === 9
@@ -29,3 +32,5 @@ o = OffsetArray(vec(A), 8);
 offset_view = @view OffsetArrays.centered(zeros(eltype(A), 5, 5))[:, begin]; # SubArray of OffsetArray
 @test @inferred(ArrayInterface.offsets(offset_view)) == (-2,)
 
+B = OffsetArray(PermutedDimsArray(rand(2,3,4), (2,3,1)));
+@test @inferred(ArrayInterface.StrideIndex(B)) === ArrayInterface.StrideIndex{3, (2, 3, 1), 3}((2, 6, static(1)), (1, 1, 1))
diff --git a/src/axes.jl b/src/axes.jl
@@ -191,9 +191,7 @@ struct LazyAxis{N,P} <: AbstractUnitRange{Int}
 end
 
 @inline Base.parent(x::LazyAxis{N,P}) where {N,P} = axes(getfield(x, :parent), static(N))
-@inline function Base.parent(x::LazyAxis{:,P}) where {P}
-    return eachindex(IndexLinear(), getfield(x, :parent))
-end
+@inline Base.parent(x::LazyAxis{:,P}) where {P} = eachindex(IndexLinear(), getfield(x, :parent))
 
 @inline parent_type(::Type{LazyAxis{N,P}}) where {N,P} = axes_types(P, static(N))
 # TODO this approach to parent_type(::Type{LazyAxis{:}}) is a bit hacky. Something like
@@ -220,7 +218,7 @@ ArrayInterfaceCore.known_first(::Type{<:LazyAxis{N,P}}) where {N,P} = known_offs
 ArrayInterfaceCore.known_first(::Type{<:LazyAxis{:,P}}) where {P} = 1
 @inline function Base.first(x::LazyAxis{N})::Int where {N}
     if ArrayInterfaceCore.known_first(x) === nothing
-        return Int(offsets(parent(x), StaticInt(N)))
+        return Int(offsets(getfield(x, :parent), StaticInt(N)))
     else
         return Int(known_first(x))
     end
@@ -229,7 +227,8 @@ end
 ArrayInterfaceCore.known_last(::Type{LazyAxis{N,P}}) where {N,P} = known_last(axes_types(P, static(N)))
 ArrayInterfaceCore.known_last(::Type{LazyAxis{:,P}}) where {P} = known_length(P)
 Base.last(x::LazyAxis) = _last(known_last(x), x)
-_last(::Nothing, x) = last(parent(x))
+_last(::Nothing, x::LazyAxis{:}) = lastindex(getfield(x, :parent))
+_last(::Nothing, x::LazyAxis{N}) where {N} = lastindex(getfield(x, :parent), N)
 _last(N::Int, x) = N
 
 known_length(::Type{<:LazyAxis{:,P}}) where {P} = known_length(P)
@@ -242,7 +241,9 @@ Base.axes1(x::LazyAxis) = x
 Base.axes(x::Slice{<:LazyAxis}) = (Base.axes1(x),)
 # assuming that lazy loaded params like dynamic length from `size(::Array, dim)` are going
 # be used again later with `Slice{LazyAxis}`, we quickly load indices
-Base.axes1(x::Slice{<:LazyAxis}) = indices(parent(x.indices))
+
+Base.axes1(x::Slice{LazyAxis{N,A}}) where {N,A} = indices(getfield(x.indices, :parent), StaticInt(N))
+Base.axes1(x::Slice{LazyAxis{:,A}}) where {A} = indices(getfield(x.indices, :parent))
 Base.to_shape(x::LazyAxis) = Base.length(x)
 
 @propagate_inbounds function Base.getindex(x::LazyAxis, i::CanonicalInt)
diff --git a/src/indexing.jl b/src/indexing.jl
@@ -183,12 +183,14 @@ to_index(::MyIndexStyle, axis, arg) = ...
 """
 to_index(x, i::Slice) = i
 to_index(x, ::Colon) = indices(x)
+to_index(::LinearIndices{0,Tuple{}}, ::Colon) = Slice(static(1):static(1))
+to_index(::CartesianIndices{0,Tuple{}}, ::Colon) = Slice(static(1):static(1))
 # logical indexing
 to_index(x, i::AbstractArray{Bool}) = LogicalIndex(i)
 to_index(x::LinearIndices, i::AbstractArray{Bool}) = LogicalIndex{Int}(i)
 # cartesian indexing
 @inline to_index(x, i::CartesianIndices{0}) = i
-@inline to_index(x, i::CartesianIndices) = axes(i)
+@inline to_index(x, i::CartesianIndices) = getfield(i, :indices)
 @inline to_index(x, i::CartesianIndex) = Tuple(i)
 @inline to_index(x, i::NDIndex) = Tuple(i)
 @inline to_index(x, i::AbstractArray{<:AbstractCartesianIndex}) = i
@@ -251,10 +253,10 @@ indices calling [`to_axis`](@ref).
     end
 end
 # drop this dimension
-to_axes(A, a::Tuple, i::Tuple{<:CanonicalInt,Vararg{Any}}) = to_axes(A, tail(a), tail(i))
+to_axes(A, a::Tuple, i::Tuple{<:CanonicalInt,Vararg{Any}}) = to_axes(A, _maybe_tail(a), tail(i))
 to_axes(A, a::Tuple, i::Tuple{I,Vararg{Any}}) where {I} = _to_axes(StaticInt(ndims_index(I)), A, a, i)
 function _to_axes(::StaticInt{1}, A, axs::Tuple, inds::Tuple)
-    return (to_axis(first(axs), first(inds)), to_axes(A, tail(axs), tail(inds))...)
+    return (to_axis(_maybe_first(axs), first(inds)), to_axes(A, _maybe_tail(axs), tail(inds))...)
 end
 @propagate_inbounds function _to_axes(::StaticInt{N}, A, axs::Tuple, inds::Tuple) where {N}
     axes_front, axes_tail = Base.IteratorsMD.split(axs, Val(N))
@@ -268,6 +270,11 @@ end
 to_axes(A, ::Tuple{Ax,Vararg{Any}}, ::Tuple{}) where {Ax} = ()
 to_axes(A, ::Tuple{}, ::Tuple{}) = ()
 
+_maybe_first(::Tuple{}) = static(1):static(1)
+_maybe_first(t::Tuple) = first(t)
+_maybe_tail(::Tuple{}) = ()
+_maybe_tail(t::Tuple) = tail(t)
+
 """
     to_axis(old_axis, index) -> new_axis
 
@@ -349,7 +356,9 @@ unsafe_getindex(A::Array, i::CanonicalInt) = Base.arrayref(false, A, Int(i))
 end
 
 unsafe_getindex(A::LinearIndices, i::CanonicalInt) = Int(i)
-unsafe_getindex(A::CartesianIndices, i::CanonicalInt, ii::Vararg{CanonicalInt}) = CartesianIndex(i, ii...)
+unsafe_getindex(A::CartesianIndices{N}, ii::Vararg{CanonicalInt,N}) where {N} = CartesianIndex(ii...)
+unsafe_getindex(A::CartesianIndices, ii::Vararg{CanonicalInt}) =
+    unsafe_getindex(A, Base.front(ii)...)
 unsafe_getindex(A::CartesianIndices, i::CanonicalInt) = @inbounds(A[i])
 
 unsafe_getindex(A::ReshapedArray, i::CanonicalInt) = @inbounds(parent(A)[i])
@@ -378,18 +387,37 @@ function unsafe_get_collection(A, inds)
 end
 _ints2range(x::CanonicalInt) = x:x
 _ints2range(x::AbstractRange) = x
+# apply _ints2range to front N elements
+_ints2range_front(::Val{N}, ind, inds...) where {N} =
+    (_ints2range(ind), _ints2range_front(Val(N - 1), inds...)...)
+_ints2range_front(::Val{0}, ind, inds...) = ()
+_ints2range_front(::Val{0}) = ()
+# get output shape with given indices
+_output_shape(::CanonicalInt, inds...) = _output_shape(inds...)
+_output_shape(ind::AbstractRange, inds...) = (length(ind), _output_shape(inds...)...)
+_output_shape(::CanonicalInt) = ()
+_output_shape(x::AbstractRange) = (length(x),)
 @inline function unsafe_get_collection(A::CartesianIndices{N}, inds) where {N}
     if (Base.length(inds) === 1 && N > 1) || stride_preserving_index(typeof(inds)) === False()
         return Base._getindex(IndexStyle(A), A, inds...)
     else
-        return CartesianIndices(to_axes(A, _ints2range.(inds)))
+        return reshape(
+            CartesianIndices(_ints2range_front(Val(N), inds...)),
+            _output_shape(inds...)
+        )
     end
 end
+_known_first_isone(ind) = known_first(ind) !== nothing && isone(known_first(ind))
 @inline function unsafe_get_collection(A::LinearIndices{N}, inds) where {N}
     if Base.length(inds) === 1 && isone(_ndims_index(typeof(inds), static(1)))
         return @inbounds(eachindex(A)[first(inds)])
-    elseif stride_preserving_index(typeof(inds)) === True()
-        return LinearIndices(to_axes(A, _ints2range.(inds)))
+    elseif stride_preserving_index(typeof(inds)) === True() &&
+            reduce_tup(&, map(_known_first_isone, inds))
+        # create a LinearIndices when first(ind) != 1 is imposable
+        return reshape(
+            LinearIndices(_ints2range_front(Val(N), inds...)),
+            _output_shape(inds...)
+        )
     else
         return Base._getindex(IndexStyle(A), A, inds...)
     end
diff --git a/src/stridelayout.jl b/src/stridelayout.jl
@@ -164,7 +164,8 @@ function contiguous_axis(::Type{T}) where {T<:PermutedDimsArray}
     end
 end
 function contiguous_axis(::Type{<:Base.ReshapedArray{T, N, A, Tuple{}}}) where {T, N, A}
-    if isone(-contiguous_axis(A))
+    c = contiguous_axis(A)
+    if c !== nothing && isone(-c)
         return StaticInt(-1)
     elseif dynamic(is_column_major(A) & is_dense(A))
         return StaticInt(1)
@@ -261,8 +262,15 @@ end
 @inline function stride_rank(::Type{A}) where {NB,NA,B<:AbstractArray{<:Any,NB},A<:Base.ReinterpretArray{<:Any,NA,<:Any,B,true}}
     NA == NB ? stride_rank(B) : _stride_rank_reinterpret(stride_rank(B), gt(StaticInt{NB}(), StaticInt{NA}()))
 end
+@inline function stride_rank(::Type{A}) where {N,B<:AbstractArray{<:Any,N},A<:Base.ReinterpretArray{<:Any,N,<:Any,B,false}}
+    stride_rank(B)
+end
+
 @inline _stride_rank_reinterpret(sr, ::False) = (One(), map(Base.Fix2(+, One()), sr)...)
 @inline _stride_rank_reinterpret(sr::Tuple{One,Vararg}, ::True) = map(Base.Fix2(-, One()), tail(sr))
+function contiguous_axis(::Type{R}) where {T,N,S,B<:AbstractArray{S,N},R<:ReinterpretArray{T,N,S,B,false}}
+    contiguous_axis(B)
+end
 # if the leading dim's `stride_rank` is not one, then that means the individual elements are split across an axis, which ArrayInterface
 # doesn't currently have a means of representing.
 @inline function contiguous_axis(::Type{A}) where {NB,NA,B<:AbstractArray{<:Any,NB},A<:Base.ReinterpretArray{<:Any,NA,<:Any,B,true}}
@@ -448,9 +456,17 @@ _dense_dims(::Type{S}, ::Nothing, ::Val{R}) where {R,N,NP,T,A<:AbstractArray{T,N
     end
 end
 
-function dense_dims(::Type{Base.ReshapedArray{T, N, P, Tuple{Vararg{Base.SignedMultiplicativeInverse{Int},M}}}}) where {T,N,P,M}
-    return _reshaped_dense_dims(dense_dims(P), is_column_major(P), Val{N}(), Val{M}())
+function dense_dims(T::Type{<:Base.ReshapedArray})
+    d = dense_dims(parent_type(T))
+    if d === nothing
+        return nothing
+    elseif all(d)
+        return n_of_x(StaticInt(ndims(T)), True())
+    else
+        return n_of_x(StaticInt(ndims(T)), False())
+    end
 end
+                
 is_dense(A) = is_dense(typeof(A))
 is_dense(::Type{A}) where {A} = _is_dense(dense_dims(A))
 _is_dense(::Tuple{False,Vararg}) = False()
@@ -459,19 +475,6 @@ _is_dense(t::Tuple{True}) = True()
 _is_dense(t::Tuple{}) = True()
 _is_dense(::Nothing) = False()
 
-
-_reshaped_dense_dims(_, __, ___, ____) = nothing
-function _reshaped_dense_dims(dense::Tuple, ::True, ::Val{N}, ::Val{0}) where {N}
-    if all(dense)
-        return _all_dense(Val{N}())
-    else
-        return nothing
-    end
-end
-function _reshaped_dense_dims(dense::Tuple{Static.False}, ::True, ::Val{N}, ::Val{0}) where {N}
-    return return ntuple(_ -> False(), Val{N}())
-end
-
 """
     known_strides(::Type{T}) -> Tuple
     known_strides(::Type{T}, dim) -> Union{Int,Nothing}
diff --git a/test/axes.jl b/test/axes.jl
@@ -11,20 +11,27 @@ m = Array{Float64}(undef, 4, 3)
 @testset "LazyAxis" begin
     A = zeros(3,4,5);
     SA = MArray(zeros(3,4,5))
+    DA = MArray(zeros(3,4,5), LinearIndices((Base.Slice(1:3), 1:4, 1:5)))
     lz1 = ArrayInterface.LazyAxis{1}(A)
     slz1 = ArrayInterface.LazyAxis{1}(SA)
+    dlz1 = ArrayInterface.LazyAxis{1}(DA)
     lzc = ArrayInterface.LazyAxis{:}(A)
     slzc = ArrayInterface.LazyAxis{:}(SA)
-
-    @test @inferred(first(lz1)) === @inferred(first(slz1))
-    @test @inferred(first(lzc)) === @inferred(first(slzc))
-    @test @inferred(last(lz1)) === @inferred(last(slz1))
-    @test @inferred(last(lzc)) === @inferred(last(slzc))
-    @test @inferred(length(lz1)) === @inferred(length(slz1))
-    @test @inferred(length(lzc)) === @inferred(length(slzc))
-    @test @inferred(Base.to_shape(lzc)) == length(slzc)
+    dlzc = ArrayInterface.LazyAxis{:}(DA)
+
+    @test @inferred(first(lz1)) === @inferred(first(slz1)) === @inferred(first(dlz1))
+    @test @inferred(first(lzc)) === @inferred(first(slzc)) === @inferred(first(dlzc))
+    @test @inferred(last(lz1)) === @inferred(last(slz1)) === @inferred(last(dlz1))
+    @test @inferred(last(lzc)) === @inferred(last(slzc)) === @inferred(last(dlzc))
+    @test @inferred(length(lz1)) === @inferred(length(slz1)) === @inferred(length(dlz1))
+    @test @inferred(length(lzc)) === @inferred(length(slzc)) === @inferred(length(dlzc))
+    @test @inferred(Base.to_shape(lzc)) == length(slzc) == length(dlzc)
     @test @inferred(Base.checkindex(Bool, lzc, 1)) & @inferred(Base.checkindex(Bool, slzc, 1))
     @test axes(lzc)[1] == Base.axes1(lzc) == axes(Base.Slice(lzc))[1] == Base.axes1(Base.Slice(lzc))
+    @test keys(axes(A, 1)) == @inferred(keys(lz1))
+
+    @test @inferred(ArrayInterface.known_first(slzc)) === 1
+    @test @inferred(ArrayInterface.known_length(slz1)) === 3
 
     @test @inferred(getindex(lz1, 2)) == 2
     @test @inferred(getindex(lz1, 1:2)) == 1:2
@@ -33,6 +40,8 @@ m = Array{Float64}(undef, 4, 3)
     @test @inferred(ArrayInterface.parent_type(ArrayInterface.LazyAxis{:}(A))) <: Base.OneTo{Int}
     @test @inferred(ArrayInterface.parent_type(ArrayInterface.LazyAxis{4}(SA))) <: ArrayInterface.SOneTo{1}
     @test @inferred(ArrayInterface.parent_type(ArrayInterface.LazyAxis{:}(SA))) <: ArrayInterface.SOneTo{60}
+    @test @inferred(IndexStyle(SA)) isa IndexLinear
+    @test @inferred(IndexStyle(DA)) isa IndexLinear
     @test ArrayInterface.can_change_size(ArrayInterface.LazyAxis{1,Vector{Any}})
 
     Aperm = PermutedDimsArray(A, (3,1,2))
diff --git a/test/indexing.jl b/test/indexing.jl
diff --git a/test/stridelayout.jl b/test/stridelayout.jl
