Add insertdims method which is inverse to dropdims

[roflmaostc]

Hi all!

My first attempt to add a method which we think is quite helpful.
It is similar to dropdims but basically the inverse to it. I tried to follow the implementation of dropdims. It also looks good in @code_warntype.

I would be happy to receive some feedback and suggestions :)

Best,

Felix

Example:

julia> a = [1 2; 3 4]
2×2 Matrix{Int64}:
 1  2
 3  4

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

[:, :, 2, 1] =
 2  4

julia> b[1,1,1,1] = 5; a
2×2 Matrix{Int64}:
 5  2
 3  4

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

[:, :, 2, 1] =
 3

[:, :, 1, 2] =
 2

[:, :, 2, 2] =
 4

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

[:, :, 1, 2] =
 2  4

[roflmaostc]

Haven't looked into the code since then.
One disadvantage of the implementation is that it only works with sorted tuples. But there is not nice way of sorting tuples at the moment so I'm not quite sure how to tackle this best

[roflmaostc]

@mkitti had a great solution which has slightly differently semantics

julia> function insertdims(A; dims)
           idx = ntuple(ndims(A) + length(dims)) do i
               if i ∈ dims
                   [CartesianIndex()]
               else
                   (:)
               end
           end
           @view A[idx...]
       end
insertdims (generic function with 1 method)

julia> A = zeros(5,4,3);

julia> size(insertdims(A; dims = (1,3,6)))
(1, 5, 1, 4, 3, 1)

julia> dropdims(b, dims=(1,2))
2×2 Matrix{Int64}:
 1  2
 3  4

julia> b = insertdims(a, dims=(1,2))
1×1×2×2 Array{Int64, 4}:
[:, :, 1, 1] =
 1

[:, :, 2, 1] =
 3

[:, :, 1, 2] =
 2

[:, :, 2, 2] =
 4

julia> dropdims(b, dims=(1,2))
2×2 Matrix{Int64}:
 1  2
 3  4

julia> b = insertdims(a, dims=(1,2))
1×1×2×2 Array{Int64, 4}:
[:, :, 1, 1] =
 1

[:, :, 2, 1] =
 3

[:, :, 1, 2] =
 2

[:, :, 2, 2] =
 4

[aplavin]

Not to discourage/criticize this PR, but: such functionality is available with a nice, consistent, and general syntax in a package:

julia> using Accessors

julia> A = zeros(5, 4, 3);

# insert singleton dimension
julia> B = @insert size(A)[1] = 1;
julia> size(B)
(1, 5, 4, 3)

# delete singleton dimension:
julia> C = @delete size(B)[1];
julia> size(C)
(5, 4, 3)

Multidim insert/delete not implemented, but can potentially be added as well.

[roflmaostc]

That looks really cool!!

I like that it is implemented but I think it still would be nice to have insertdims since dropdims is already existing.

[aplavin]

That looks really cool!!

Yeah, and the best part - it's all totally composable!

julia> A = zeros(5, 4, 3, 1);

# drop singleton dimension, wherever it is:
julia> B = @delete size(A) |> _[findfirst(==(1), _)];
julia> size(B)
(5, 4, 3)

# insert singleton dimension last:
julia> C = @insert last(size(A)) = 1);
julia> size(C)
(5, 4, 3, 1, 1)

[mkitti]

@mkitti had a great solution which has slightly differently semantics

julia> function insertdims(A; dims)
           idx = ntuple(ndims(A) + length(dims)) do i
               if i ∈ dims
                   [CartesianIndex()]
               else
                   (:)
               end
           end
           @view A[idx...]
       end
insertdims (generic function with 1 method)

One thing I like about the semantics of my implementation is that the following is mostly true:

julia> A = rand(3,4,5); dims = (1,4,6);

julia> dropdims(insertdims(A; dims); dims) == A
true

There might be a flaw though:

julia> size(insertdims(A; dims=7))
(3, 4, 5, 1)

Perhaps we need the following

julia> function insertdims(A; dims)
           idx = ntuple(maximum((ndims(A) + length(dims), dims...))) do i
               if i ∈ dims
                   [CartesianIndex()]
               else
                   (:)
               end
           end
           @view A[idx...]
       end
insertdims (generic function with 1 method)

julia> size(insertdims(A; dims=7))
(3, 4, 5, 1, 1, 1, 1)

[mkitti]

From triage, WWMBD: what would @mbauman do?

The traditional answer is just spell insertdims in the following way:

const newaxis = [CartesianIndex()]
@view array[:, newaxis, :, ...]

An alternative to this pull request would be document the above better and/or define the above constant. Has any thinking changed?

[roflmaostc]

Perhaps we need the following

julia> function insertdims(A; dims)
           idx = ntuple(maximum((ndims(A) + length(dims), dims...))) do i
               if i ∈ dims
                   [CartesianIndex()]
               else
                   (:)
               end
           end
           @view A[idx...]
       end
insertdims (generic function with 1 method)

julia> size(insertdims(A; dims=7))
(3, 4, 5, 1, 1, 1, 1)

That's not type stable, is it?

[mkitti]

Perhaps we need the following

julia> function insertdims(A; dims)
           idx = ntuple(maximum((ndims(A) + length(dims), dims...))) do i
               if i ∈ dims
                   [CartesianIndex()]
               else
                   (:)
               end
           end
           @view A[idx...]
       end
insertdims (generic function with 1 method)

julia> size(insertdims(A; dims=7))
(3, 4, 5, 1, 1, 1, 1)

That's not type stable, is it?

We could consider a type stable version in some circumstances, such as providing dims as a Val.

[mbauman]

I think doing this as a reshape operation as initially proposed here makes far more sense than an indexing one. The reshape is going to be significantly easier on both type stability and preserving strided-ness... and it really is more reshape-y than index-y.

The only reason to use indexing operations is to match the surface of numpy's np.newaxis idioms. But I think insertdims is a better API in any case — and then its implementation can just be whatever makes the most sense.

[roflmaostc]

The only issue in my implementation was that it required a sorted input for the inserted dims.
Is there a way to sort tuples with Julia Bases in the meantime?

[LilithHafner]

Is there a way to sort tuples with Julia Bases in the meantime?

If someone (could it be you?) reviews #54494 there will be

[roflmaostc]

If @mbauman thinks that requiring sorted tuples as input is fine, we could lift that requirement also later.

[roflmaostc]

In NumPy the convention is different:

>>> x = np.array([[1, 2], [3,4]])
>>> np.expand_dims(x, axis=(0, 1,3)).shape
(1, 1, 2, 1, 2)

Personally I find this a bit confusing since it counts the inserted dims too.

[mkitti]

Numpy seems to match the semantics I proposed earlier. I understand this in terms of where I want the singleton dimensions to be in the final result.

[nsajko]

The doc string should be included into the docs:

diff --git a/doc/src/base/arrays.md b/doc/src/base/arrays.md
index ccfb237..66fe5c7 100644
--- a/doc/src/base/arrays.md
+++ b/doc/src/base/arrays.md
@@ -138,6 +138,7 @@ Base.parentindices
 Base.selectdim
 Base.reinterpret
 Base.reshape
+Base.insertdims
 Base.dropdims
 Base.vec
 Base.SubArray

[roflmaostc]

The way it works right now, it inserts singleton dimensions at the position of the dims indicated with respect to the initial input array dimensions.

Let's say Array shape is (2,3,4) and I do dims=(1,2,3). It would insert a singleton dimension before the 2, before the 3 and before the 4.

julia> size(_insertdims(ones(2,3,4), (1,2,3)))
(1, 2, 1, 3, 1, 4)

I am not sure if there is a semantic which is easier to comprehend but in my experience we often had a 2D array and want it to be 3D or ND.
In this case you could simply call: insertdims(A, dims=(3,))

But I am totally open for suggestions if we can find something else

[roflmaostc]

It also claims to be the inverse of dropdims but looking at these examples it clearly isn't.

I thought in the following way: insertdims insert singleton dimensions, dropdims removes them.

The other way of doing it, would be

julia> size(dropdims(ones((1,1, 2, 3, 1, 4)), dims=(1,2, 5)))
(2, 3, 4)

julia> insertdims2(ones((2,3,4)), dims=(1,2,5))
(1, 1, 2, 3, 1, 4)

For me this requires some acrobatics, where the dimension with size 4 ends up, because I need to insert the dims virtually in my head. Would it be before or after the singleton dimension?
On the other hand with my convention, it's immediately clear where the singleton dimensions are (with respect to the input dimensions).

Does it make sense?

[aplavin]

The way it works right now, it inserts singleton dimensions at the position of the dims indicated with respect to the initial input array dimensions.

Hmm, I see... That definitely wasn't my first intuition :)
Maybe, if keeping this semantics, it's cleaner to say "add singleton dimensions before dimensions number dims in the original array" or something...

But also, maybe the "inverse to dropdims" semantics is cleaner? Whenever possible.
This corresponds to interpreting dims as indices in the resulting array dimensions.
Also, nicely meshes with potential extensions by various "named arrays": dims=:mynewdim clearly corresponds to the resulting array.

[roflmaostc]

@LilithHafner thanks for your suggestions. To be consistent with permutedims, I changed it.

I also included your docstring and your tests.

[nsajko]

The type system doesn't ensure N isa Int, but we can be explicit about this to help Julia's abstract type inference.

[LilithHafner]

Triage is happy with adding this new function and exporting it because it mirrors dropdims.

[roflmaostc]

ok 🆒
Let me know if I can do more.

[LilithHafner]

Thank you @roflmaostc! I appreciate your commitment and follow-through.