How to drop singleton dimensions in a type stable way?

General Usage
1

I have a 3D array A which is obtained by calculating a certain function over a grid specified by axes xs, ys and zs. These axes can be either an AbstractVector{Number} or a Number. When any of the axes is a Number, the corresponding dimension has size 1. Is there a concise way to write a type stable function drop_singletons(A, xs, ys, zs) that drops all the singleton dimensions whose corresponding axis is a Number?

For example, when

xs = 1:2
ys = 1
zs = 1:3

I would have size(A) = (2,1,3) and I would like that

size(drop_singletons(A, xs, ys, zs)) = (2,3)

I could do it case by case for all possible signatures, e.g, by defining

function drop_singletons(A, x, ::Number, z)
    reshape(A, length(x), length(z))
end

and so on, but that is very cumbersome and prone to mistakes.

Can anybody see a better way? Thanks for the help!

2

Base.dropdims? Not certain what circumstances you’re doing this so I can’t guarantee type stability, but it’ll specialize on the keyword argument dims (I only ever use Int or a tuple of Int for it).

3 
Base.@assume_effects :foldable function drop_singleton_dims(@nospecialize dims::Tuple)
    if dims === ()
        dims
    else
        let f = dims[1], r = Base.tail(dims), z = @inline drop_singleton_dims(r)
            z = z::Tuple
            if f isa AbstractVector
                (f, z...)::Tuple{Any,Vararg}
            else
                z
            end
        end
    end::Tuple
end

function drop_singletons(arr, dims::Vararg{Any,N}) where {N}
    d = drop_singleton_dims(dims)
    lengths = map(length, d)
    reshape(arr, lengths)
end
1 Like
4

Here’s another alternative:

maybedropdims(A; dims=0) = A
maybedropdims(A, ::Number, args...; dims=1) = maybedropdims(dropdims(A; dims), args...; dims)
maybedropdims(A, ::AbstractVector, args...; dims=1) = maybedropdims(A, args...; dims=dims+1)

julia> A = randn(3,1,5,2);

julia> maybedropdims(A, 1:3, 1, 1:5, 1:2) |> typeof
Array{Float64, 3}

Seems to run equally fast as a straight dropdims(A; dims=2) call.

Edit: Fixed a bug. Also, this is less efficient when there are multiple dimensions to drop.

5

This is the native behavior of getindex. Vector slices yield output dimensions but dimensions indexed by scalars are dropped. You should just be able to do

julia> A = zeros(2,1,3);

julia> A[xs,ys,zs] # `ys` is scalar so is dropped from the output dimensions
2×3 Matrix{Float64}:
 0.0  0.0  0.0
 0.0  0.0  0.0
2 Likes
6

Eh. I now came up with a solution that matches dropdims, but just using getindex is easier (though my new solution is actually a bit faster.) But if you use view you get almost zero cost:

view(A, xs, ys, zs)

7

Nice, but this copies.

Nice, but this can’t return the same type (usually it’ll return SubArray).

The solutions based on getindex or view are very short, but an ideal solution basically must be based on reshape.

8

First, let’s focus on the function countvecs below.

julia> function countvecs(args::Union{AbstractVector{<: Number}, Number}...)
           count(args) do arg
               arg isa AbstractVector
           end
       end
countvecs (generic function with 1 method)
julia> code_llvm(countvecs, Tuple{UnitRange{Int}, Int, Int})
;  @ REPL[18]:1 within `countvecs`
define i64 @julia_countvecs_892([2 x i64]* nocapture readonly %0, i64 signext %1, i64 signext %2) #0 {
top:
  ret i64 1
}

julia> code_llvm(countvecs, Tuple{UnitRange{Int}, Vector{Int}, Int})
;  @ REPL[18]:1 within `countvecs`
define i64 @julia_countvecs_894([2 x i64]* nocapture readonly %0, {}* %1, i64 signext %2) #0 {
top:
;  @ REPL[18]:2 within `countvecs`
  ret i64 2
}

julia> code_llvm(countvecs, Tuple{UnitRange{Int}, Vector{Int}, Vector{Int}})
;  @ REPL[18]:1 within `countvecs`
define i64 @julia_countvecs_896([2 x i64]* nocapture readonly %0, {}* %1, {}* %2) #0 {
top:
;  @ REPL[18]:2 within `countvecs`
  ret i64 3
}

We can see that given arguments of concrete types, the function is constant. This gives is the dimensionality of the output array, allowing for complete type stability.

julia> function drop_singletons(
           A::Array{T},
           args...
       ) where {T}
           N = countvecs(args...)
           sz = NTuple{N,Int}(
               [length(a) for a in args if !(a isa Number)]
           )
           reshape(A, sz)::Array{T,N}
       end
drop_singletons (generic function with 1 methods)

julia> @code_warntype drop_singletons(zeros(2,1,3), 1:2, 5, 1:3)
MethodInstance for drop_singletons(::Array{Float64, 3}, ::UnitRange{Int64}, ::Int64, ::UnitRange{Int64})
  from drop_singletons(A::Array{T}, args...) where T @ Main REPL[57]:1
Static Parameters
  T = Float64
Arguments
  #self#::Core.Const(drop_singletons)
  A::Array{Float64, 3}
  args::Tuple{UnitRange{Int64}, Int64, UnitRange{Int64}}
Locals
  #31::var"#31#33"
  #30::var"#30#32"
  sz::Tuple{Int64, Int64}
  N::Int64
Body::Matrix{Float64}
1 ─       (N = Core._apply_iterate(Base.iterate, Main.countvecs, args))
│   %2  = Core.apply_type(Main.NTuple, N::Core.Const(2), Main.Int)::Core.Const(Tuple{Int64, Int64})
│         (#30 = %new(Main.:(var"#30#32")))
│   %4  = #30::Core.Const(var"#30#32"())
│         (#31 = %new(Main.:(var"#31#33")))
│   %6  = #31::Core.Const(var"#31#33"())
│   %7  = Base.Filter(%6, args)::Base.Iterators.Filter{var"#31#33", Tuple{UnitRange{Int64}, Int64, UnitRange{Int64}}}
│   %8  = Base.Generator(%4, %7)::Base.Generator{Base.Iterators.Filter{var"#31#33", Tuple{UnitRange{Int64}, Int64, UnitRange{Int64}}}, var"#30#32"}
│   %9  = Base.collect(%8)::Vector{Int64}
│         (sz = (%2)(%9))
│   %11 = Main.reshape(A, sz)::Matrix{Float64}
│   %12 = Core.apply_type(Main.Array, $(Expr(:static_parameter, 1)), N::Core.Const(2))::Core.Const(Matrix{Float64})
│   %13 = Core.typeassert(%11, %12)::Matrix{Float64}
└──       return %13

I do want to point out that this will not drop axes such as [1] or 1:1.

9

In that case, here’s my solution. It seems to match @nsajko’s, but uses a different approach:

getsize() = ()
getsize(::Number, args...) = getsize(args...)
getsize(x::AbstractVector, args...) = (length(x), getsize(args...)...)
maybereshape(A, args...) = reshape(A, getsize(args...))

julia> B = randn(3,1,5,1);

julia> @btime maybereshape($B, 1:3, 1, 1:5, 1)
  18.136 ns (1 allocation: 64 bytes)
3×5 Matrix{Float64}:
  0.268776   0.0475721   0.206281  -1.03224   -0.275928
 -0.326243   0.602314   -0.511709  -0.527623  -1.43498
 -0.837188  -0.261605    0.403192  -1.13667   -0.496302

julia> @btime getsize(1:3, 1, 1:5, 1)
  1.100 ns (0 allocations: 0 bytes)
(3, 5)
2 Likes