Skip to content

Commit c7344b4

Browse files
jishnubjishnub
committed
Fix ambiguities and preserve Integer types
1 parent c87b8f8 commit c7344b4

File tree

4 files changed

+185
-32
lines changed

4 files changed

+185
-32
lines changed

‎base/abstractarray.jl

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -835,13 +835,11 @@ similar(a::AbstractArray, ::Type{T}, dims::Dims{N}) where {T,N} = Array{T,N}(
835835
to_shape(::Tuple{}) = ()
836836
to_shape(dims::Dims) = dims
837837
to_shape(dims::DimsOrInds) = map(to_shape, dims)::DimsOrInds
838-
to_shape(dims::Tuple{Vararg{Union{Integer, AbstractUnitRange, Colon}}}) = map(to_shape, dims)
839838
# each dimension
840839
to_shape(i::Int) = i
841840
to_shape(i::Integer) = Int(i)
842841
to_shape(r::OneTo) = Int(last(r))
843842
to_shape(r::AbstractUnitRange) = r
844-
to_shape(r::Colon) = r
845843

846844
"""
847845
similar(storagetype, axes)

‎base/reshapedarray.jl

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -121,22 +121,31 @@ reshape
121121

122122
# we collect the vararg indices and only define methods for tuples of indices
123123
reshape(parent::AbstractArray, dims::Union{Integer,Colon,AbstractUnitRange}...) = reshape(parent, dims)
124-
reshape(parent::AbstractArray, dims::Tuple{Vararg{Integer}}) = _reshape(parent, map(Int, dims))
124+
reshape(parent::AbstractArray, dims::Tuple{Vararg{Integer}}) = reshape(parent, map(Int, dims))
125+
# Usually, custom arrays would only need to specialize `reshape(A::MyArray, dims::Tuple{Vararg{Int}})`, and other
126+
# integer types passed as `dims` would be converted to `Int`s
127+
reshape(parent::AbstractArray, dims::Dims) = _reshape(parent, dims)
125128

126129
# Allow missing dimensions with Colon():
127-
# convert axes to sizes using to_shape, and convert colons to sizes using _reshape_uncolon
130+
# Replace `OneTo`s by their lengths, and convert colons to sizes using _reshape_uncolon
128131
# We add a level of indirection to avoid method ambiguities in reshape
129132
reshape(parent::AbstractArray, dims::Tuple{Vararg{Union{Integer,Colon,OneTo}}}) = _reshape_maybecolon(parent, dims)
130133
_reshape_maybecolon(parent::AbstractVector, ::Tuple{Colon}) = parent
131-
_reshape_maybecolon(parent::AbstractArray, dims::Tuple{Vararg{Union{Integer,Colon,OneTo}}}) = reshape(parent, _reshape_uncolon(length(parent), to_shape(dims)))
134+
# helper function analogous to to_shape, but this doesn't cast Integers to Int
135+
# this allows dispatching to a specialized reshape(M::MyArray, dims::Tuple{Vararg{Integer}})
136+
_to_shape(i::OneTo) = length(i)
137+
_to_shape(i::Union{Integer, Colon}) = i
138+
_to_shape(t::Tuple) = map(_to_shape, t)
139+
_reshape_maybecolon(parent::AbstractArray, dims::Tuple{Vararg{Union{Integer,Colon,OneTo}}}) = reshape(parent, _reshape_uncolon(parent, _to_shape(dims)))
132140

133141
@noinline _reshape_throwcolon(dims) = throw(DimensionMismatch(LazyString("new dimensions ", dims,
134142
" may have at most one omitted dimension specified by `Colon()`")))
135143
@noinline _reshape_throwsize(lenA, dims) = throw(DimensionMismatch(LazyString("array size ", lenA,
136144
" must be divisible by the product of the new dimensions ", dims)))
137145

138-
_reshape_uncolon(len, ::Tuple{Colon}) = len
139-
@inline function _reshape_uncolon(len, _dims::Tuple{Vararg{Union{Integer, Colon}}})
146+
_reshape_uncolon(parent::AbstractArray, dims::Tuple) = _reshape_uncolon(length(parent), dims)
147+
_reshape_uncolon(len::Integer, ::Tuple{Colon}) = len
148+
@inline function _reshape_uncolon(len::Integer, _dims::Tuple{Vararg{Union{Integer, Colon}}})
140149
# promote the dims to `Int` at least
141150
dims = map(x -> x isa Colon ? x : promote_type(typeof(x), Int)(x), _dims)
142151
dims isa Tuple{Vararg{Integer}} && return dims

‎test/offsetarray.jl

Lines changed: 171 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,23 @@
22

33
isdefined(Main, :OffsetArrays) || @eval Main include("testhelpers/OffsetArrays.jl")
44
using .Main.OffsetArrays
5-
import .Main.OffsetArrays: IdOffsetRange
5+
import .Main.OffsetArrays: IdOffsetRange, no_offset_view
66
using Random
77
using LinearAlgebra
88
using Base: IdentityUnitRange
99
using Test
1010

11+
no_offset_axes(x, d) = no_offset_view(axes(x, d))
12+
no_offset_axes(x) = map(no_offset_view, axes(x))
13+
14+
function same_value(r1, r2)
15+
length(r1) == length(r2) || return false
16+
for (v1, v2) in zip(r1, r2)
17+
v1 == v2 || return false
18+
end
19+
return true
20+
end
21+
1122
if !isdefined(@__MODULE__, :T24Linear)
1223
include("testhelpers/arrayindexingtypes.jl")
1324
end
@@ -846,6 +857,84 @@ end
846857
a = OffsetArray(4:5, 5:6)
847858
@test reshape(a, :) === a
848859
@test reshape(a, (:,)) === a
860+
861+
A0 = [1 3; 2 4]
862+
A = OffsetArray(A0, (-1,2))
863+
864+
B = reshape(A0, -10:-9, 9:10)
865+
@test isa(B, OffsetArray{Int,2})
866+
@test parent(B) == A0
867+
@test axes(B) == IdentityUnitRange.((-10:-9, 9:10))
868+
B = reshape(A, -10:-9, 9:10)
869+
@test isa(B, OffsetArray{Int,2})
870+
@test axes(B) == IdentityUnitRange.((-10:-9, 9:10))
871+
b = reshape(A, -7:-4)
872+
@test axes(b) == (IdentityUnitRange(-7:-4),)
873+
@test isa(parent(b), Vector{Int})
874+
@test parent(b) == A0[:]
875+
a = OffsetArray(rand(3,3,3), -1:1, 0:2, 3:5)
876+
# Offset axes are required for reshape(::OffsetArray, ::Val) support
877+
b = reshape(a, Val(2))
878+
@test isa(b, OffsetArray{Float64,2})
879+
@test axes(b) == IdentityUnitRange.((-1:1, 1:9))
880+
b = reshape(a, Val(4))
881+
@test isa(b, OffsetArray{Float64,4})
882+
@test axes(b) == (axes(a)..., IdentityUnitRange(1:1))
883+
884+
@test reshape(OffsetArray(-1:0, -1:0), :, 1) == reshape(-1:0, 2, 1)
885+
@test reshape(OffsetArray(-1:2, -1:2), -2:-1, :) == reshape(-1:2, -2:-1, 2)
886+
887+
@test reshape(OffsetArray(-1:0, -1:0), :) == OffsetArray(-1:0, -1:0)
888+
@test reshape(A, :) == reshape(A0, :)
889+
890+
# reshape with one Colon for AbstractArrays
891+
B = reshape(A0, -10:-9, :)
892+
@test B isa OffsetArray{Int,2}
893+
@test parent(B) == A0
894+
@test no_offset_axes(B, 1) == -10:-9
895+
@test axes(B, 2) == axes(A0, 2)
896+
897+
B = reshape(A0, -10:-9, 3:3, :)
898+
@test B isa OffsetArray{Int,3}
899+
@test same_value(A0, B)
900+
@test no_offset_axes(B, 1) == -10:-9
901+
@test no_offset_axes(B, 2) == 3:3
902+
@test axes(B, 3) == 1:2
903+
904+
B = reshape(A0, -10:-9, 3:4, :)
905+
@test B isa OffsetArray{Int,3}
906+
@test same_value(A0, B)
907+
@test no_offset_axes(B, 1) == -10:-9
908+
@test no_offset_axes(B, 2) == 3:4
909+
@test axes(B, 3) == 1:1
910+
911+
# pop the parent
912+
B = reshape(A, size(A))
913+
@test B == A0
914+
B = reshape(A, (Base.OneTo(2), 2))
915+
@test B == A0
916+
B = reshape(A, (2,:))
917+
@test B == A0
918+
919+
# julialang/julia #33614
920+
A = OffsetArray(-1:0, (-2,))
921+
@test reshape(A, :) == A
922+
Arsc = reshape(A, :, 1)
923+
Arss = reshape(A, 2, 1)
924+
@test Arsc[1,1] == Arss[1,1] == -1
925+
@test Arsc[2,1] == Arss[2,1] == 0
926+
@test_throws BoundsError Arsc[0,1]
927+
@test_throws BoundsError Arss[0,1]
928+
A = OffsetArray([-1,0], (-2,))
929+
Arsc = reshape(A, :, 1)
930+
Arsc[1,1] = 5
931+
@test first(A) == 5
932+
933+
Vec64 = zeros(6)
934+
ind_a_64 = 3
935+
ind_a_32 =Int32.(ind_a_64)
936+
@test reshape(Vec64, ind_a_32, :) == reshape(Vec64, ind_a_64, :)
937+
849938
R = reshape(zeros(6), 2, :)
850939
@test R isa Matrix
851940
@test axes(R) == (1:2, 1:3)
@@ -856,6 +945,87 @@ end
856945
@test axes(R) == (2:3, 1:3)
857946
R = reshape(zeros(6,1), 2:3, :)
858947
@test axes(R) == (2:3, 1:3)
948+
949+
R = reshape(zeros(6), 1:2, :)
950+
@test axes(R) == (1:2, 1:3)
951+
R = reshape(zeros(6,1), 1:2, :)
952+
@test axes(R) == (1:2, 1:3)
953+
954+
# reshape works even if the parent doesn't have 1-based indices
955+
# this works even if the parent doesn't support the reshape
956+
r = OffsetArray(IdentityUnitRange(0:1), -1)
957+
@test reshape(r, 2) == 0:1
958+
@test reshape(r, (2,)) == 0:1
959+
@test reshape(r, :) == OffsetArray(0:1, -1:0)
960+
@test reshape(r, (:,)) == OffsetArray(0:1, -1:0)
961+
962+
@test reshape(ones(2:3, 4:5), (2, :)) == ones(2,2)
963+
964+
# more than one colon is not allowed
965+
@test_throws Exception reshape(ones(3:4, 4:5, 1:2), :, :, 2)
966+
@test_throws Exception reshape(ones(3:4, 4:5, 1:2), :, 2, :)
967+
968+
A = OffsetArray(rand(4, 4), -1, -1);
969+
B = reshape(A, (2, :))
970+
@test axes(B, 1) == 1:2
971+
@test axes(B, 2) == 1:8
972+
973+
# reshape with one single colon becomes a `vec`
974+
A = OffsetArray(rand(4, 4), -1, -1)
975+
@test reshape(A, (:, )) == vec(A)
976+
@test reshape(A, :) == vec(A)
977+
978+
A0 = [1 3; 2 4]
979+
A = reshape(A0, 2:3, 4:5)
980+
@test axes(A) == Base.IdentityUnitRange.((2:3, 4:5))
981+
@test reshape(A, axes(parent(A))..., :) == reshape(A0, axes(A0)..., :)
982+
983+
B = reshape(A0, -10:-9, 9:10)
984+
@test isa(B, OffsetArray{Int,2})
985+
@test parent(B) == A0
986+
@test axes(B) == Base.IdentityUnitRange.((-10:-9, 9:10))
987+
988+
@testset "BigInt reshape" begin
989+
struct MyBigFill{T,N} <: AbstractArray{T,N}
990+
val :: T
991+
axes :: NTuple{N,Base.OneTo{BigInt}}
992+
end
993+
MyBigFill(val, sz::Tuple{}) = MyBigFill{typeof(val),0}(val, sz)
994+
MyBigFill(val, sz::Tuple{Vararg{Integer}}) = MyBigFill(val, map(Base.OneTo{BigInt}, sz))
995+
Base.size(M::MyBigFill) = map(length, M.axes)
996+
Base.axes(M::MyBigFill) = M.axes
997+
function Base.getindex(M::MyBigFill{<:Any,N}, ind::Vararg{Integer,N}) where {N}
998+
checkbounds(M, ind...)
999+
M.val
1000+
end
1001+
function Base.isassigned(M::MyBigFill{<:Any,N}, ind::Vararg{BigInt,N}) where {N}
1002+
checkbounds(M, ind...)
1003+
true
1004+
end
1005+
function Base.reshape(M::MyBigFill, ind::NTuple{N,BigInt}) where {N}
1006+
length(M) == prod(ind) || throw(ArgumentError("length mismatch in reshape"))
1007+
MyBigFill(M.val, ind)
1008+
end
1009+
Base.reshape(M::MyBigFill, ind::Tuple{}) = MyBigFill(M.val, ind)
1010+
1011+
M = MyBigFill(4, (2, 3))
1012+
O = OffsetArray(M)
1013+
@test vec(O) isa MyBigFill
1014+
@test vec(O) == vec(M)
1015+
@test reshape(O, :, big(2)) == reshape(M, :, big(2))
1016+
@test reshape(O, :, big(2)) isa MyBigFill
1017+
@test reshape(O, axes(M)) isa MyBigFill
1018+
@test reshape(O, axes(M)) == reshape(M, axes(M))
1019+
@test reshape(O, axes(M,1), :) isa MyBigFill
1020+
@test reshape(O, axes(M,1), :) == reshape(M, axes(M,1), :)
1021+
@test reshape(O, axes(M,1), big(1), :) isa MyBigFill
1022+
@test reshape(O, axes(M,1), big(1), :) == reshape(M, axes(M,1), big(1), :)
1023+
1024+
M = MyBigFill(4, (1,1))
1025+
O = OffsetArray(M)
1026+
@test reshape(O) isa MyBigFill
1027+
@test reshape(O) == reshape(M)
1028+
end
8591029
end
8601030

8611031
@testset "stack" begin
@@ -924,14 +1094,3 @@ end
9241094
b = sum(a, dims=1)
9251095
@test b[begin] == sum(r)
9261096
end
927-
928-
@testset "reshape" begin
929-
A0 = [1 3; 2 4]
930-
A = reshape(A0, 2:3, 4:5)
931-
@test axes(A) == Base.IdentityUnitRange.((2:3, 4:5))
932-
933-
B = reshape(A0, -10:-9, 9:10)
934-
@test isa(B, OffsetArray{Int,2})
935-
@test parent(B) == A0
936-
@test axes(B) == Base.IdentityUnitRange.((-10:-9, 9:10))
937-
end

‎test/testhelpers/OffsetArrays.jl

Lines changed: 0 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -552,22 +552,9 @@ _reshape2(A, inds) = reshape(A, inds)
552552
_reshape2(A::OffsetArray, inds) = reshape(parent(A), inds)
553553
_reshape_nov(A, inds) = _reshape(no_offset_view(A), inds)
554554

555-
Base.reshape(A::OffsetArray, inds::Tuple{OffsetAxis,Vararg{OffsetAxis}}) =
556-
OffsetArray(_reshape(parent(A), inds), map(_toaxis, inds))
557555
# And for non-offset axes, we can just return a reshape of the parent directly
558-
Base.reshape(A::OffsetArray, inds::Tuple{Union{Integer,Base.OneTo},Vararg{Union{Integer,Base.OneTo}}}) = _reshape_nov(A, inds)
559556
Base.reshape(A::OffsetArray, inds::Tuple{Integer,Vararg{Integer}}) = _reshape_nov(A, inds)
560-
Base.reshape(A::OffsetArray, inds::Tuple{Union{Colon, Integer}, Vararg{Union{Colon, Integer}}}) = _reshape_nov(A, inds)
561557
Base.reshape(A::OffsetArray, inds::Dims) = _reshape_nov(A, inds)
562-
Base.reshape(A::OffsetVector, ::Colon) = A
563-
Base.reshape(A::OffsetVector, ::Tuple{Colon}) = A
564-
Base.reshape(A::OffsetArray, inds::Union{Int,Colon}...) = reshape(A, inds)
565-
Base.reshape(A::OffsetArray, inds::Tuple{Vararg{Union{Int,Colon}}}) = _reshape_nov(A, inds)
566-
# The following two additional methods for Colon are added to resolve method ambiguities to
567-
# Base: https://github.com/JuliaLang/julia/pull/45387#issuecomment-1132859663
568-
Base.reshape(A::OffsetArray, inds::Colon) = _reshape_nov(A, inds)
569-
Base.reshape(A::OffsetArray, inds::Tuple{Colon}) = _reshape_nov(A, inds)
570-
571558
# permutedims in Base does not preserve axes, and can not be fixed in a non-breaking way
572559
# This is a stopgap solution
573560
Base.permutedims(v::OffsetVector) = reshape(v, (1, axes(v, 1)))

0 commit comments

Comments
 (0)