Add squeeze(A) method to squeeze all singleton dimensions

JuliaLang · May 21, 2017 · b91d0eb · b91d0eb
1 parent f988566
commit b91d0eb
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 0 deletions.
diff --git a/base/abstractarraymath.jl b/base/abstractarraymath.jl
@@ -81,6 +81,29 @@ end
 
 squeeze(A::AbstractArray, dim::Integer) = squeeze(A, (Int(dim),))
 
+"""
+    squeeze(A)
+
+Remove all singleton dimensions from array `A`.
+
+```jldoctest
+julia> a = reshape(collect(1:4),(2,2,1,1))
+2×2×1×1 Array{Int64,4}:
+[:, :, 1, 1] =
+ 1  3
+ 2  4
+
+julia> squeeze(a)
+2×2 Array{Int64,2}:
+ 1  3
+ 2  4
+```
+"""
+function squeeze(A::AbstractArray)
+    singleton_dims = [d for d in 1:ndims(A) if size(A, d) == 1]
+    return squeeze(A, tuple(singleton_dims...))
+end
+
 
 ## Unary operators ##
 

diff --git a/test/subarray.jl b/test/subarray.jl
@@ -555,3 +555,11 @@ let
     @test Base.IndexStyle(view(a, :, :)) == Base.IndexLinear()
     @test isbits(view(a, :, :))
 end
+
+# issue
+# Ensure that we can auto-squeeze out all singleton dimensions of an array
+let
+    @test ndims(squeeze(ones(10, 1, 1))) == 1
+    @test ndims(squeeze(ones(1, 10))) == 1
+    @test ndims(squeeze(ones(10, 10))) == 2
+end