Combining `eachindex` and `axes`?

One of the things I like about eachindex is that you can give it several arrays and it throws a DimensionMismatch if they have different sizes.

I wonder if it would make sense to enable the same with axes(A, d), possibly choosing a different d for each A. With such a syntax, I would write a matrix-vector product like this:

function matvec!(y::Vector{T}, A::Matrix, x::Vector) where {T}
    for i in axes((y, A), (1, 1))
        y[i] = zero(T)
        for j in axes((A, x), (2, 1))
            y[i] += A[i, j] * x[j]
        end
    end
end

and the indexing would be safe by default, which means I could sprinkle @inbounds and @simd. Currently, this is of course feasible but a little more clumsy:

function matvec!(y::Vector{T}, A::Matrix, x::Vector) where {T}
    @assert axes(y, 1) == axes(A, 1)
    @assert axes(x, 1) == axes(A, 2)
    for i in axes(y, 1)
        y[i] = zero(T)
        for j in axes(x, 1)
            y[i] += A[i, j] * x[j]
        end
    end
end
3 Likes

StaticArrayInteface.jl has an indices function which does exactly this:

julia> using StaticArrayInterface: indices

julia> function matvec!(y::Vector{T}, A::Matrix, x::Vector) where {T}
           for i in indices((y, A), (1, 1))
               y[i] = zero(T)
               for j in indices((A, x), (2, 1))
                   y[i] += A[i, j] * x[j]
               end
           end
       end;

julia> let x = rand(10), A = rand(10, 10), y = similar(x)
           matvec!(y, A, x)
           y ≈ A * x
       end
true
2 Likes

Good to know that it exists! I guess I was wondering if it deserves a place in Base to encourage correctness when iterating on several arrays

1 Like