Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add AbstractOneTo and have OneTo be its subtype #56902

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
Open

Add AbstractOneTo and have OneTo be its subtype #56902

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
34a8f77
Add AbstractOneTo ang have OneTo be its subtype
jishnub Dec 21, 2024
3742f4f
Pass axes/dims to `similar` in 2-arg methods
jishnub Dec 25, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 10 additions & 5 deletions base/abstractarray.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -818,15 +818,18 @@ julia> similar(falses(10), Float64, 2, 4)
See also: [`undef`](@ref), [`isassigned`](@ref).
"""
similar(a::AbstractArray{T}) where {T} = similar(a, T)
similar(a::AbstractArray, ::Type{T}) where {T} = similar(a, T, to_shape(axes(a)))
similar(a::AbstractArray{T}, dims::Tuple) where {T} = similar(a, T, to_shape(dims))
similar(a::AbstractArray{T}, dims::DimOrInd...) where {T} = similar(a, T, to_shape(dims))
similar(a::AbstractArray, ::Type{T}, dims::DimOrInd...) where {T} = similar(a, T, to_shape(dims))
similar(a::AbstractArray, ::Type{T}) where {T} = similar(a, T, axes(a))
similar(a::AbstractArray{T}, dims::Tuple) where {T} = similar(a, T, dims)
similar(a::AbstractArray{T}, dims::DimOrInd...) where {T} = similar(a, T, dims)
similar(a::AbstractArray, ::Type{T}, dims::DimOrInd...) where {T} = similar(a, T, dims)
# Similar supports specifying dims as either Integers or AbstractUnitRanges or any mixed combination
# thereof. Ideally, we'd just convert Integers to OneTos and then call a canonical method with the axes,
# but we don't want to require all AbstractArray subtypes to dispatch on Base.OneTo. So instead we
# define this method to convert supported axes to Ints, with the expectation that an offset array
# package will define a method with dims::Tuple{Union{Integer, UnitRange}, Vararg{Union{Integer, UnitRange}}}
similar(a::AbstractArray, ::Type{T}, dims::Tuple{Union{Integer, AbstractOneTo}, Vararg{Union{Integer, AbstractOneTo}}}) where {T} = similar(a, T, to_shape(dims))
# legacy method for packages that specialize similar(A::AbstractArray, ::Type{T}, dims::Tuple{Union{Integer, OneTo, CustomAxis}, Vararg{Union{Integer, OneTo, CustomAxis}}}
# leaving this method in ensures that Base owns the more specific method
similar(a::AbstractArray, ::Type{T}, dims::Tuple{Union{Integer, OneTo}, Vararg{Union{Integer, OneTo}}}) where {T} = similar(a, T, to_shape(dims))
# similar creates an Array by default
similar(a::AbstractArray, ::Type{T}, dims::Dims{N}) where {T,N} = Array{T,N}(undef, dims)
Expand All @@ -837,7 +840,7 @@ to_shape(dims::DimsOrInds) = map(to_shape, dims)::DimsOrInds
# each dimension
to_shape(i::Int) = i
to_shape(i::Integer) = Int(i)
to_shape(r::OneTo) = Int(last(r))
to_shape(r::AbstractOneTo) = Int(last(r))
to_shape(r::AbstractUnitRange) = r

"""
Expand All @@ -863,6 +866,8 @@ would create a 1-dimensional logical array whose indices match those
of the columns of `A`.
"""
similar(::Type{T}, dims::DimOrInd...) where {T<:AbstractArray} = similar(T, dims)
similar(::Type{T}, shape::Tuple{Union{Integer, AbstractOneTo}, Vararg{Union{Integer, AbstractOneTo}}}) where {T<:AbstractArray} = similar(T, to_shape(shape))
# legacy method for packages that specialize similar(::Type{T}, dims::Tuple{Union{Integer, OneTo, CustomAxis}, Vararg{Union{Integer, OneTo, CustomAxis}})
similar(::Type{T}, shape::Tuple{Union{Integer, OneTo}, Vararg{Union{Integer, OneTo}}}) where {T<:AbstractArray} = similar(T, to_shape(shape))
similar(::Type{T}, dims::Dims) where {T<:AbstractArray} = T(undef, dims)

Expand Down
9 changes: 8 additions & 1 deletion base/range.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -451,14 +451,21 @@ if isdefined(Main, :Base)
end
end

"""
Base.AbstractOneTo

Abstract type for ranges that start at 1.
"""
abstract type AbstractOneTo{T} <: AbstractUnitRange{T} end

"""
Base.OneTo(n)

Define an `AbstractUnitRange` that behaves like `1:n`, with the added
distinction that the lower limit is guaranteed (by the type system) to
be 1.
"""
struct OneTo{T<:Integer} <: AbstractUnitRange{T}
struct OneTo{T<:Integer} <: AbstractOneTo{T}
stop::T # invariant: stop >= zero(stop)
function OneTo{T}(stop) where {T<:Integer}
throwbool(r) = (@noinline; throw(ArgumentError("invalid index: $r of type Bool")))
Expand Down
3 changes: 3 additions & 0 deletions base/reshapedarray.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -120,6 +120,9 @@ julia> reshape(1:6, 2, 3)
reshape

reshape(parent::AbstractArray, dims::IntOrInd...) = reshape(parent, dims)
reshape(parent::AbstractArray, shp::Tuple{Union{Integer,AbstractOneTo}, Vararg{Union{Integer,AbstractOneTo}}}) = reshape(parent, to_shape(shp))
# legacy method for packages that specialize reshape(parent::AbstractArray, shp::Tuple{Union{Integer,OneTo,CustomAxis}, Vararg{Union{Integer,OneTo,CustomAxis}}})
# leaving this method in ensures that Base owns the more specific method
reshape(parent::AbstractArray, shp::Tuple{Union{Integer,OneTo}, Vararg{Union{Integer,OneTo}}}) = reshape(parent, to_shape(shp))
reshape(parent::AbstractArray, dims::Tuple{Integer, Vararg{Integer}}) = reshape(parent, map(Int, dims))
reshape(parent::AbstractArray, dims::Dims) = _reshape(parent, dims)
Expand Down
21 changes: 21 additions & 0 deletions test/abstractarray.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2227,3 +2227,24 @@ end
@test_throws MethodError isreal(G)
end
end

@testset "similar/reshape for AbstractOneTo" begin
A = [1,2]
@testset "reshape" begin
@test reshape(A, 2, SizedArrays.SOneTo(1)) == reshape(A, 2, 1)
@test reshape(A, Base.OneTo(2), SizedArrays.SOneTo(1)) == reshape(A, 2, 1)
@test reshape(A, SizedArrays.SOneTo(1), 2) == reshape(A, 1, 2)
@test reshape(A, SizedArrays.SOneTo(1), Base.OneTo(2)) == reshape(A, 1, 2)
end
@testset "similar" begin
b = similar(A, SizedArrays.SOneTo(1), big(2))
@test b isa Array{Int, 2}
@test size(b) == (1, 2)
b = similar(A, SizedArrays.SOneTo(1), Base.OneTo(2))
@test b isa Array{Int, 2}
@test size(b) == (1, 2)
b = similar(A, SizedArrays.SOneTo(1), 2, Base.OneTo(2))
@test b isa Array{Int, 3}
@test size(b) == (1, 2, 2)
end
end
14 changes: 5 additions & 9 deletions test/testhelpers/SizedArrays.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ import LinearAlgebra: mul!

export SizedArray

struct SOneTo{N} <: AbstractUnitRange{Int} end
struct SOneTo{N} <: Base.AbstractOneTo{Int} end
SOneTo(N) = SOneTo{N}()
Base.length(::SOneTo{N}) where {N} = N
Base.size(r::SOneTo) = (length(r),)
Expand Down Expand Up @@ -58,14 +58,6 @@ Base.parent(S::SizedArray) = S.data
+(S1::SizedArray{SZ}, S2::SizedArray{SZ}) where {SZ} = SizedArray{SZ}(S1.data + S2.data)
==(S1::SizedArray{SZ}, S2::SizedArray{SZ}) where {SZ} = S1.data == S2.data

homogenize_shape(t::Tuple) = (_homogenize_shape(first(t)), homogenize_shape(Base.tail(t))...)
homogenize_shape(::Tuple{}) = ()
_homogenize_shape(x::Integer) = x
_homogenize_shape(x::AbstractUnitRange) = length(x)
const Dims = Union{Integer, Base.OneTo, SOneTo}
function Base.similar(::Type{A}, shape::Tuple{Dims, Vararg{Dims}}) where {A<:AbstractArray}
similar(A, homogenize_shape(shape))
end
function Base.similar(::Type{A}, shape::Tuple{SOneTo, Vararg{SOneTo}}) where {A<:AbstractArray}
R = similar(A, length.(shape))
SizedArray{length.(shape)}(R)
Expand All @@ -74,6 +66,10 @@ function Base.similar(x::SizedArray, ::Type{T}, shape::Tuple{SOneTo, Vararg{SOne
sz = map(length, shape)
SizedArray{sz}(similar(parent(x), T, sz))
end
function Base.reshape(x::AbstractArray, shape::Tuple{SOneTo, Vararg{SOneTo}})
sz = map(length, shape)
SizedArray{length.(sz)}(reshape(x, length.(sz)))
end

const SizedMatrixLike = Union{SizedMatrix, Transpose{<:Any, <:SizedMatrix}, Adjoint{<:Any, <:SizedMatrix}}

Expand Down
Loading