Repository: JuliaLang/julia
File: base/abstractarraymath.jl (Julia)

Code before:

julia> b[1,1,1] = 5; a
2×2×1×1 Array{Int64, 4}:
[:, :, 1, 1] =
 5  3
 2  4
```
"""
dropdims(A; dims) = _dropdims(A, dims)
function _dropdims(A::AbstractArray, dims::Dims)
    for i in eachindex(dims)
        1 <= dims[i] <= ndims(A) || throw(ArgumentError("dropped dims must be in range 1:ndims(A)"))
        length(axes(A, dims[i])) == 1 || throw(ArgumentError("dropped dims must all be size 1"))
        for j = 1:i-1
            dims[j] == dims[i] && throw(ArgumentError("dropped dims must be unique"))
        end
    end
    ax = _foldoneto((ds, d) -> d in dims ? ds : (ds..., axes(A,d)), (), Val(ndims(A)))
    reshape(A, ax::typeof(_sub(axes(A), dims)))
end
_dropdims(A::AbstractArray, dim::Integer) = _dropdims(A, (Int(dim),))


"""
    insertdims(A; dims)

Inverse of [`dropdims`](@ref); return an array with new singleton dimensions
at every dimension in `dims`.

Repeated dimensions are forbidden and the largest entry in `dims` must be
less than or equal than `ndims(A) + length(dims)`.

The result shares the same underlying data as `A`, such that the
result is mutable if and only if `A` is mutable, and setting elements of one
alters the values of the other.

See also: [`dropdims`](@ref), [`reshape`](@ref), [`vec`](@ref).
# Examples
```jldoctest
julia> x = [1 2 3; 4 5 6]
2×3 Matrix{Int64}:
 1  2  3
 4  5  6

julia> insertdims(x, dims=3)
2×3×1 Array{Int64, 3}:
[:, :, 1] =
 1  2  3
 4  5  6

julia> insertdims(x, dims=(1,2,5)) == reshape(x, 1, 1, 2, 3, 1)

Review comment (refactor):
I'd deduplicate this by moving the `reshape` call out of the branch (assuming the `typeassert` is necessary in the first place):

```suggestion
    if isconcretetype(eltype(ox))
        # if all the axes are the same type, we can use the tail as the
        # axes of the result rather than extracting one at each index
        ax = ax::typeof(_sub(ox, dims))
    end
    return reshape(A, ax)
```

Code after:

julia> b[1,1,1] = 5; a
2×2×1×1 Array{Int64, 4}:
[:, :, 1, 1] =
 5  3
 2  4
```
"""
dropdims(A; dims) = _dropdims(A, dims)
function _dropdims(A::AbstractArray, dims::Dims)
    for i in eachindex(dims)
        1 <= dims[i] <= ndims(A) || throw(ArgumentError("dropped dims must be in range 1:ndims(A)"))
        length(axes(A, dims[i])) == 1 || throw(ArgumentError("dropped dims must all be size 1"))
        for j = 1:i-1
            dims[j] == dims[i] && throw(ArgumentError("dropped dims must be unique"))
        end
    end
    ox = axes(A)
    ax = _foldoneto((ds, d) -> d in dims ? ds : (ds..., axes(A,d)), (), Val(ndims(A)))
    if isconcretetype(eltype(ox))
        # if all the axes are the same type, we can use the tail as the
        # axes of the result rather than extracting one at each index
        return reshape(A, ax::typeof(_sub(ox, dims)))
    else
        return reshape(A, ax)
    end
end
_dropdims(A::AbstractArray, dim::Integer) = _dropdims(A, (Int(dim),))

"""
    insertdims(A; dims)

Inverse of [`dropdims`](@ref); return an array with new singleton dimensions
at every dimension in `dims`.

Repeated dimensions are forbidden and the largest entry in `dims` must be
less than or equal than `ndims(A) + length(dims)`.

The result shares the same underlying data as `A`, such that the
result is mutable if and only if `A` is mutable, and setting elements of one
alters the values of the other.

See also: [`dropdims`](@ref), [`reshape`](@ref), [`vec`](@ref).
# Examples
```jldoctest
julia> x = [1 2 3; 4 5 6]
2×3 Matrix{Int64}:
 1  2  3
 4  5  6

julia> insertdims(x, dims=3)