Safe co-iteration along axes

[Tokazama]

This came up with the following comment in from #61 .

Users will often check for equal sizes in front of a loop, and eachindex will guarantee this as well

for I in eachindex(A,C) # checks for equivalent size
   
end

Because eachindex isn't always appropriate (e.g., when you don't have exact correspondence between axes of the arrays), I think it'd be a good to have something like an eachindexalongaxis:

function mymatmulkernel!(C, A, B)
    @avx for m in eachindexalongaxis((C,A),(1,1)), n in eachindexalongaxis((C,B),(2,2)), k in eachindexalongaxis((A,B),(2,1))
        C[m,n] += A[m,k] * B[k,n]
    end
end

where it checks sizes. eachindexalongaxis or each_index_along_axis doesn't really roll off > the tongue, and I don't know where such a function should live.
I'd also have to think a little more about how to use it and actually pass the information of > dimension equality (to be combined with density guarantees) to deduce stride-equality. The hacky way would be to just pattern match eachindex and eachindexalongindex in the macro, handling dimension equality through the expression, while handling density info based on types.

I'd like to propose the use of indices for specifically retrieving the iterator that corresponds to each axis of an array. It would function like axes but if there's any sort of other data associated with an axis it throws it away so there's a clear relationship between memory layout and and each dimension.

In practice it would look like this:

julia> indices(ones(2,3))
(Base.OneTo(2), Base.OneTo(3))

julia> indices(ones(2,3), 1)
Base.OneTo(2)

julia> indices(ones(2,3), 2)
Base.OneTo(3)

julia> indices(ones(2,3), ones(2,3), 2)
Base.OneTo(3)

julia> indices(ones(2,3), ones(2,2), 2)
ERROR: DimensionMismatch("all inputs to indices must have the same indices, got Base.OneTo(3) and Base.OneTo(2)")

In this case the first several examples are the same as if using axes, but it helps for cases like those in AxisIndices.jl where each axis can have additional information (e.g., keys) associated with the indices. The last two examples would ensure that the iterator for each collection is the same along the second dimension.

[chriselrod]

In my example I used a tuple of indices, so that we could

julia> indices(ones(2,3), ones(2,2), (1,1))
Base.OneTo(2)

julia> indices(ones(2,3), ones(2,2), (1,2))
Base.OneTo(2)

julia> indices(ones(2,3), ones(2,2), (2,1))
ERROR: DimensionMismatch("all inputs to indices must have the same indices, got Base.OneTo(3) and Base.OneTo(2)")

julia> indices(ones(2,3), ones(2,2), (2,2))
ERROR: DimensionMismatch("all inputs to indices must have the same indices, got Base.OneTo(3) and Base.OneTo(2)")

As we may wish to iterate across different axes for the different arrays. My example also used two tuples, as it's easier to define a function in this way:

indices(A::AbstractArray, i::Integer) = axes(A, i)
@inline index_tuple(i::Integer, n) = i # either index into `i`, or return it, depending on whether it's a tuple or scalar
@inline index_tuple(i::Tuple{Vararg{<:Integer}}, n) = i[n]
function indices(A::Tuple{Vararg{<:AbstractArray,N}}, i) where {N}
    inds = ntuple(n -> axes(A[n], index_tuple(i, n)), Val(N))
    ind = first(inds)
    @assert all(isequal(ind), Base.tail(inds)) "Not all specified axes are equal: $inds"
    ind
end

This should have more or less the desired behavior.

julia> indices((ones(2,3), ones(2,2)), 1)
 Base.OneTo(2)

julia> indices((ones(2,3), ones(2,2)), 2)
ERROR: AssertionError: Not all specified axes are equal: (Base.OneTo(3), Base.OneTo(2))
Stacktrace:
 [1] indices(A::Tuple{Matrix{Float64},Matrix{Float64}}, i::Int64)
   @ Main ./REPL[228]:4
 [2] top-level scope
   @ REPL[238]:1

To edit a specific method, type the corresponding number into the REPL and press Ctrl+Q

julia> indices((ones(2,3), ones(2,2)), (1,1))
 Base.OneTo(2)

julia> indices((ones(2,3), ones(2,2)), (1,2))
 Base.OneTo(2)

julia> indices((ones(2,3), ones(2,2)), (2,1))
ERROR: AssertionError: Not all specified axes are equal: (Base.OneTo(3), Base.OneTo(2))
Stacktrace:
 [1] indices(A::Tuple{Matrix{Float64},Matrix{Float64}}, i::Tuple{Int64,Int64})
   @ Main ./REPL[228]:4
 [2] top-level scope
   @ REPL[235]:1

To edit a specific method, type the corresponding number into the REPL and press Ctrl+Q

julia> indices((ones(2,3), ones(2,2)), (2,2))
ERROR: AssertionError: Not all specified axes are equal: (Base.OneTo(3), Base.OneTo(2))
Stacktrace:
 [1] indices(A::Tuple{Matrix{Float64},Matrix{Float64}}, i::Tuple{Int64,Int64})
   @ Main ./REPL[228]:4
 [2] top-level scope
   @ REPL[236]:1

To edit a specific method, type the corresponding number into the REPL and press Ctrl+Q

[Tokazama]

As we may wish to iterate across different axes for the different arrays.

That's a really good point.

I've actually really wanted this feature, so if people like this I can make a PR.

[chriselrod]

Sounds good.
As an obvious improvement to my definition, it should try to return the most specific axis information available, e.g. favoring Base.OneTo over UnitRange, rather than just returning the first axis.
There's probably a better way than this:

pick_range(r1, r2) = r1
pick_range(r1::Base.OneTo, ::Any) = r1
pick_range(::Any, r2::Base.OneTo) = r2
pick_range(r1::Base.OneTo, ::Base.OneTo) = r1
reduce(pick_range, (1:5, Base.OneTo(5)))
reduce(pick_range, (Base.OneTo(5), 1:5))

A good version would be easily extensible without requiring a tremendous amount of methods to specify relative priorities, so that it could also favor StaticRanges, for example. Using ArrayInterface:

using ArrayInterface
function known_score(::R) where {R}
    kf = ArrayInterface.known_first(R)
    kl = ArrayInterface.known_last(R)
    ks = ArrayInterface.known_step(R)
    # not knowing step is much worse, so it is weighted more heavily
    isnothing(ks)*8 + isnothing(kf) + isnothing(kl)
end
function pick_range(r1::R1, r2::R2) where {R1,R2}
    r1score = known_score(r1)
    r2score = known_score(r2)
    r1score > r2score ? r2 : r1
end

Of course, a better solution would be to combine all the known information.
Hypothetically, maybe there are three axis that each provide one compile-time known piece of the puzzle.
This function should be able to consolidate all those missing pieces and return a fully static range type.

There aren't any fully static ranges in base, but we could define a bare-bones implementation here.

struct OptionallyStaticRange{F,L,S} <: AbstractRange
    first::F
    last::L
    step::S
end
ArrayInterface.known_first(::OptionallyStaticRange{Val{F}}) where {F} = F
ArrayInterface.known_last(::OptionallyStaticRange{<:Any,Val{L}}) where {L} = L
ArrayInterface.known_step(::OptionallyStaticRange{<:Any,<:Any,Val{S}}) where {S} =S
# add methods to support ArrayInterface

If they aren't Val types, then it'll return nothing, indicating unknown.

We could then have our indices function construct one of these OptionallyStaticRanges by combining information from each of the axis.