Skip to content

Commit

Permalink
Support using colon to omit a dimension in reshape
Browse files Browse the repository at this point in the history
Closes #16790.
  • Loading branch information
mbauman committed Jan 10, 2017
1 parent 28a11ff commit c5a444e
Show file tree
Hide file tree
Showing 3 changed files with 90 additions and 33 deletions.
33 changes: 0 additions & 33 deletions base/docs/helpdb/Base.jl
Original file line number Diff line number Diff line change
Expand Up @@ -446,39 +446,6 @@ See also [`zeros`](@ref), [`similar`](@ref).
"""
ones

"""
reshape(A, dims)
Create an array with the same data as the given array, but with different dimensions.
```jldoctest
julia> A = collect(1:16)
16-element Array{Int64,1}:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
julia> reshape(A, (2, 8))
2×8 Array{Int64,2}:
1 3 5 7 9 11 13 15
2 4 6 8 10 12 14 16
```
"""
reshape

"""
randsubseq!(S, A, p)
Expand Down
69 changes: 69 additions & 0 deletions base/reshapedarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -36,10 +36,79 @@ start(R::ReshapedArrayIterator) = start(R.iter)
end
length(R::ReshapedArrayIterator) = length(R.iter)

"""
reshape(A, dims...)
reshape(A, dims)
Return an array with the same data as the given array, but with different dimensions.
The new dimensions may be specified either as a list of arguments or as a shape
tuple. At most one dimension may be specified with a `:`, in which case its
length is computed such that its product with all the specified dimensions is
equal to the length of the original array A.
```jldoctest
julia> A = collect(1:16)
16-element Array{Int64,1}:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
julia> reshape(A, (4, 4))
4×4 Array{Int64,2}:
1 5 9 13
2 6 10 14
3 7 11 15
4 8 12 16
julia> reshape(A, 2, :)
2×8 Array{Int64,2}:
1 3 5 7 9 11 13 15
2 4 6 8 10 12 14 16
```
"""
reshape

reshape(parent::AbstractArray, dims::IntOrInd...) = reshape(parent, dims)
reshape(parent::AbstractArray, shp::NeedsShaping) = reshape(parent, to_shape(shp))
reshape(parent::AbstractArray, dims::Dims) = _reshape(parent, dims)

# Allow missing dimensions with Colon():
reshape(parent::AbstractArray, dims::Int...) = reshape(parent, dims)
reshape(parent::AbstractArray, dims::Union{Int,Colon}...) = reshape(parent, dims)
reshape(parent::AbstractArray, dims::Tuple{Vararg{Union{Int,Colon}}}) = _reshape(parent, _reshape_uncolon(parent, dims))
# Recursively move dimensions to pre and post tuples, splitting on the Colon
@inline _reshape_uncolon(A, dims) = _reshape_uncolon(A, (), nothing, (), dims)
@inline _reshape_uncolon(A, pre, c::Void, post, dims::Tuple{Any, Vararg{Any}}) =
_reshape_uncolon(A, (pre..., dims[1]), c, post, tail(dims))
@inline _reshape_uncolon(A, pre, c::Void, post, dims::Tuple{Colon, Vararg{Any}}) =
_reshape_uncolon(A, pre, dims[1], post, tail(dims))
@inline _reshape_uncolon(A, pre, c::Colon, post, dims::Tuple{Any, Vararg{Any}}) =
_reshape_uncolon(A, pre, c, (post..., dims[1]), tail(dims))
_reshape_uncolon(A, pre, c::Colon, post, dims::Tuple{Colon, Vararg{Any}}) =
throw(DimensionMismatch("new dimensions $((pre..., c, post..., dims...)) may only have at most one omitted dimension specified by Colon()"))
@inline function _reshape_uncolon(A, pre, c::Colon, post, dims::Tuple{})
sz, remainder = divrem(length(A), prod(pre)*prod(post))
remainder == 0 || _throw_reshape_colon_dimmismatch(A, pre, post)
(pre..., sz, post...)
end
_throw_reshape_colon_dimmismatch(A, pre, post) =
throw(DimensionMismatch("array size $(length(A)) must be divisible by the product of the new dimensions $((pre..., :, post...))"))

reshape{T,N}(parent::AbstractArray{T,N}, ndims::Type{Val{N}}) = parent
function reshape{T,AN,N}(parent::AbstractArray{T,AN}, ndims::Type{Val{N}})
reshape(parent, rdims((), indices(parent), Val{N}))
Expand Down
21 changes: 21 additions & 0 deletions test/arrayops.jl
Original file line number Diff line number Diff line change
Expand Up @@ -135,6 +135,27 @@ end
@test isa(reshape(s, Val{N}), Base.ReshapedArray{Int,N})
end
end
@testset "reshape with colon" begin
# Reshape with an omitted dimension
let A = linspace(1, 60, 60)
@test size(reshape(A, :)) == (60,)
@test size(reshape(A, :, 1)) == (60, 1)
@test size(reshape(A, (:, 2))) == (30, 2)
@test size(reshape(A, 3, :)) == (3, 20)
@test size(reshape(A, 2, 3, :)) == (2, 3, 10)
@test size(reshape(A, (2, :, 5))) == (2, 6, 5)
@test_throws DimensionMismatch reshape(A, 7, :)
@test_throws DimensionMismatch reshape(A, :, 2, 3, 4)
@test_throws DimensionMismatch reshape(A, (:, :))

B = rand(2,2,2,2)
@test size(reshape(B, :)) == (16,)
@test size(reshape(B, :, 4)) == (4, 4)
@test size(reshape(B, (2, 1, :))) == (2, 1, 8)
@test_throws DimensionMismatch reshape(B, 3, :)
@test_throws DimensionMismatch reshape(B, :, :, 2, 2)
end
end

@test reshape(1:5, (5,)) === 1:5
@test reshape(1:5, 5) === 1:5
Expand Down

0 comments on commit c5a444e

Please sign in to comment.