JuliaArrays / ArrayInterface.jl

Designs for new Base array interface primitives, used widely through scientific machine learning (SciML) and other organizations
https://docs.sciml.ai/ArrayInterface/stable/
MIT License
134 stars 37 forks source link

Fix `getindex` with additional inds #312

Closed wangl-cc closed 2 years ago

wangl-cc commented 2 years ago

Cuurently, getindex(A, inds...) with inds whose length is larger than ndims(A) will cause a error, but it is allowed for Base.getindex:

julia> A = reshape(1:12, 3, 4)
3×4 reshape(::UnitRange{Int64}, 3, 4) with eltype Int64:
 1  4  7  10
 2  5  8  11
 3  6  9  12

julia> ArrayInterface.getindex(A, 1, 1, :)
ERROR: ArgumentError: tuple must be non-empty

julia> A[1, 1, :]
1-element Vector{Int64}:
 1

The reason is to_axes only works length(inds) == ndims(A) currently, This PR fixes it with _maybe_first and _maybe_tail.

However, the behavior for CartesianIndices is different, which works currently but ignores additional inds

julia> ArrayInterface.getindex(CartesianIndices(A), 1, 1, :)
0-dimensional Array{CartesianIndex{2}, 0}:
CartesianIndex(1, 1)

julia> ArrayInterface.getindex(CartesianIndices(A), 1, 1, :, :)
0-dimensional Array{CartesianIndex{2}, 0}:
CartesianIndex(1, 1)

julia> ArrayInterface.getindex(CartesianIndices(A), 1, :, :, :)
4-element Vector{CartesianIndex{2}}:
 CartesianIndex(1, 1)
 CartesianIndex(1, 2)
 CartesianIndex(1, 3)
 CartesianIndex(1, 4)

julia> CartesianIndices(A)[1, :, :, :]
4×1×1 Array{CartesianIndex{2}, 3}:
[:, :, 1] =
 CartesianIndex(1, 1)
 CartesianIndex(1, 2)
 CartesianIndex(1, 3)
 CartesianIndex(1, 4)

How can I fix it or ignore it?

codecov[bot] commented 2 years ago

Codecov Report

Merging #312 (3e324ae) into master (d9b5089) will increase coverage by 0.29%. The diff coverage is 100.00%.

@@            Coverage Diff             @@
##           master     #312      +/-   ##
==========================================
+ Coverage   91.49%   91.79%   +0.29%     
==========================================
  Files           9        9              
  Lines        1399     1413      +14     
==========================================
+ Hits         1280     1297      +17     
+ Misses        119      116       -3     
Impacted Files Coverage Δ
src/indexing.jl 89.09% <100.00%> (+2.19%) :arrow_up:

Continue to review full report at Codecov.

Legend - Click here to learn more Δ = absolute <relative> (impact), ø = not affected, ? = missing data Powered by Codecov. Last update d9b5089...3e324ae. Read the comment docs.

ChrisRackauckas commented 2 years ago

Reshape to get the right sized output?

wangl-cc commented 2 years ago

Thanks for your suggestion @ChrisRackauckas, but I found that the main reason of the problem is that the additional : are converted to a zero dimensional CartesianIndices, and then ignored by Base._getindex (in below function). The similar problems occurs for LinearIndices.

https://github.com/JuliaArrays/ArrayInterface.jl/blob/d9b50899b419fe606255d9172a5e4b55171113fc/src/indexing.jl#L381-L387

But if I defined

to_index(::LinearIndices{0,Tuple{}}, ::Colon) = Slice(static(1):static(1))
to_index(::CartesianIndices{0,Tuple{}}, ::Colon) = Slice(static(1):static(1))

the stride_preserving_index(typeof(inds)) will be True(), and CartesianIndices(to_axes(A, _ints2range.(inds))) doesn't return excepted shape. Then I found that ArrayInterface.getindex(::CartesianIndices, I...) works very different from Base.getindex even for normal indices:

julia> ArrayInterface.getindex(CartesianIndices((3, 3)), 1, 1, 1)
CartesianIndex(1, 1, 1)

julia> Base.getindex(CartesianIndices((3, 3)), 1, 1, 1)
CartesianIndex(1, 1)

julia> ArrayInterface.getindex(CartesianIndices((3, 3)), 1, :)
1×3 CartesianIndices{2, Tuple{ArrayInterface.OptionallyStaticUnitRange{StaticInt{1}, Int64}, Base.OneTo{Int64}}} with indices 1:1:1×Base.OneTo(3):
 CartesianIndex(1, 1)  CartesianIndex(1, 2)  CartesianIndex(1, 3)

julia> Base.getindex(CartesianIndices((3, 3)), 1, :)
3-element Vector{CartesianIndex{2}}:
 CartesianIndex(1, 1)
 CartesianIndex(1, 2)
 CartesianIndex(1, 3)

I'm not sure if it is designed in purpose or not?

wangl-cc commented 2 years ago

There is a commit 4334830a41a7c61f6488aaa61537e82af458e93d which make ArrayInterface.getindex(::CartesianIndices, I...) works the same as Base. But I'm not sure if it can be merged.

Tokazama commented 2 years ago

A lot of this was designed with the goal to avoid allocating new arrays whenever possible, but there's been a lot of work over the years on improving CartesianIndices in base (some of which has come as a result of discussions in ArrayInterface). I'd be open to changing things here, especially if it didn't break any of the current tests.

wangl-cc commented 2 years ago

I revert to construct a CartesianIndices when possible and reshape it to correct shape. But for LinearIndices it seems impossible if first of indices is not one.

Besides, construct with to_axes(A, inds) will cause the similar issue to #308, so I remove it.

Tokazama commented 2 years ago

What we really need is our own Indices type that we can optimize however we want

Tokazama commented 2 years ago

@wangl-cc , it looks like all tests are passing and the changes seem reasonable to me. There's certainly more we can do related to this but that may be a lot for a single PR. Let me know if you're comfortable with the state of this PR and I can merge it. I'd of course be happy to review more

wangl-cc commented 2 years ago

I may not have enough time on this PR this week, so merge this PR now is okay for me. If it's possible, I can do more related works on this weekend, maybe a new PR? @Tokazama