Dropdims but an order of magnitude faster

Internals & Design
,
1

The function dropdims uses the short but (intermediately) type-instable pattern

ax = _foldoneto((ds, d) -> d in dims ? ds : (ds..., axes(A,d)), (), Val(ndims(A)))
reshape(A, ax::typeof(_sub(axes(A), dims)))

Replacing this part with dropinds below, or basically anything reasonable avoiding length-of-tuple type instability, saves up to an order magnitude in time as well as allocations. As far as I can see it is equally compilation-heavy.

My guess is that dropdims has simply never been high priority, but maybe I am missing something.

If the point is code brevity, working explicitly with a vector instead is still better (see further below). Probably some kind of case destinction concerning N would be best.

function dropinds(::Val{N}, inds::NTuple{M,Int}) where {N,M}
    @assert allunique(inds) "inds = $inds does not have unique entries"
    @assert all(x->x in 1:N, inds) "inds = $inds not within 1:$N"
    _dropinds(Val(N), inds, 0)
end
function _dropinds(::Val{N}, inds::NTuple{M,Int}, k::Int) where {N,M}
    N == M && return tuple()
    while N-k in inds
        k = k+1
    end
    return (_dropinds(Val(N-1), inds, k+1)..., N-k)
end
function dropdims_by_dropinds(x::Array{T,N}, dims::Tuple{Vararg{Int}}) where {T,N}
    t = dropinds(Val(N), dims)
    new_s = map(i->axes(x, i), t)
    reshape(x, new_s)
end

function test()
    a = randn(1, 2, 1, 1, 2, 1, 1)
    dims = Tuple(findall(==(1), size(a)))
    s = 0.0
    @time for i = 1:1000; s += dropdims_by_dropinds(a, dims)[1]; end
    @time for i = 1:1000; s += dropdims(a; dims)[1]; end

    a = randn(1, 2, 1, 1, 2, 1, 1, 1, 1, 1, 1, 2, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1)
    dims = Tuple(findall(==(1), size(a)))
    @time for i = 1:1000; s += dropdims_by_dropinds(a, dims)[1]; end
    @time for i = 1:1000; s += dropdims(a; dims)[1]; end
    return s
end
test()
# 0.000066 seconds (3.00 k allocations: 78.125 KiB)
# 0.000366 seconds (10.00 k allocations: 281.250 KiB)
# 0.000187 seconds (3.00 k allocations: 93.750 KiB)
# 0.001141 seconds (28.00 k allocations: 890.625 KiB)

Vector version:

function dropinds_vec(::Val{N}, inds::NTuple{M,Int}) where {N,M}
    @assert allunique(inds) "inds = $inds does not have unique entries"
    @assert all(x->x in 1:N, inds) "inds = $inds not within 1:$N"
    
    s = collect(1:N)
    for x = inds
        s[x] = N+1
    end
    w = 1
    for v = s
        s[w] = v
        w = ifelse(v == N+1, w, w+1)
    end  
    return ntuple(i->s[i], N-M)
end
function dropdims_by_dropinds_vec(x::Array{T,N}, dims::Tuple{Vararg{Int}}) where {T,N}
    t = dropinds_vec(Val(N), dims)
    new_s = map(i->axes(x, i), t)
    reshape(x, new_s)
end
function test()
    a = randn(1, 2, 1, 1, 2, 1, 1)
    dims = Tuple(findall(==(1), size(a)))
    s = 0.0
    @time for i = 1:1000; s += dropdims_by_dropinds_vec(a, dims)[1]; end
    @time for i = 1:1000; s += dropdims(a; dims)[1]; end

    a = randn(1, 2, 1, 1, 2, 1, 1, 1, 1, 1, 1, 2, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1)
    dims = Tuple(findall(==(1), size(a)))
    @time for i = 1:1000; s += dropdims_by_dropinds_vec(a, dims)[1]; end
    @time for i = 1:1000; s += dropdims(a; dims)[1]; end
    return s
end
test()
# 0.000089 seconds (5.00 k allocations: 187.500 KiB)
# 0.000508 seconds (10.00 k allocations: 281.250 KiB)
# 0.000188 seconds (5.00 k allocations: 328.125 KiB)
# 0.001286 seconds (28.00 k allocations: 890.625 KiB)
4 Likes