Add dropdims(f, args..; dims, kwargs..) for reductions to drop dims

NEWS.md

@@ -26,14 +26,14 @@ New library functions
 * The `tempname` function now takes a `cleanup::Bool` keyword argument defaulting to `true`, which causes the process to try to ensure that any file or directory at the path returned by `tempname` is deleted upon process exit ([#33090]).
 * The `readdir` function now takes a `join::Bool` keyword argument defaulting to `false`, which when set causes `readdir` to join its directory argument with each listed name ([#33113]).
 
-
 Standard library changes
 ------------------------
 
 * The methods of `mktemp` and `mktempdir` which take a function body to pass temporary paths to no longer throw errors if the path is already deleted when the function body returns ([#33091]).
 
-#### Libdl
+* A new `dropdims(f, args...; dims, kwargs...)` method computes the reduction `f` over the region described by `dims` and then drops those dimensions from the result ([#33130]).
 
+#### Libdl
 
 #### LinearAlgebra
 

base/abstractarraymath.jl

@@ -50,7 +50,7 @@ _sub(t::Tuple, s::Tuple) = _sub(tail(t), tail(s))
 
 Remove the dimensions specified by `dims` from array `A`.
 Elements of `dims` must be unique and within the range `1:ndims(A)`.
-`size(A,i)` must equal 1 for all `i` in `dims`.
+`size(A, i)` must equal 1 for all `i` in `dims`.
 
 # Examples
 ```jldoctest
@@ -65,6 +65,11 @@ julia> dropdims(a; dims=3)
 [:, :, 1] =
  1  3
  2  4
+
+julia> dropdims(a; dims=(3, 4))
+2×2 Array{Int64,2}:
+ 1  3
+ 2  4
 ```
 """
 dropdims(A; dims) = _dropdims(A, dims)
@@ -84,8 +89,42 @@ function _dropdims(A::AbstractArray, dims::Dims)
     end
     reshape(A, d::typeof(_sub(axes(A), dims)))
 end
+
 _dropdims(A::AbstractArray, dim::Integer) = _dropdims(A, (Int(dim),))
 
+"""
+    dropdims(f, args...; dims, kwargs...)
+
+Compute reduction `f` over dimensions `dims` and drop those dimensions from the result.
+
+The reduction `f` must both accept a `dims` keyword argument, and reduce the the dimensions
+specified by `dims` to size 1.
+
+!!! compat "Julia 1.4"
+    This method requires at least Julia 1.4.
+
+# Examples
+```jldoctest
+julia> a = [3.0  2.0  6.0  8.0
+            6.0  1.0  4.0  2.0
+            3.0  0.0  7.0  6.0];
+
+julia> dropdims(sum, a, dims=1)
+4-element Array{Float64,1}:
+ 12.0
+  3.0
+ 17.0
+ 16.0
+
+julia> dropdims(sum, abs2, a, dims=2)
+3-element Array{Float64,1}:
+ 113.0
+  57.0
+  94.0
+```
+"""
+dropdims(f, args...; dims, kwargs...) = _dropdims(f(args...; kwargs..., dims=dims), dims)
+
 ## Unary operators ##
 
 conj(x::AbstractArray{<:Real}) = x

test/arrayops.jl

@@ -303,6 +303,37 @@ end
     @test_throws ArgumentError dropdims(a, dims=4)
     @test_throws ArgumentError dropdims(a, dims=6)
 
+    # dropdims with reductions. issue #16606
+    @test (@inferred(dropdims(sum, a, dims=1)) ==
+           @inferred(dropdims(sum, a, dims=(1,))) ==
+           reshape(sum(a, dims=1), (1, 8, 8, 1)))
+    @test (@inferred(dropdims(sum, a, dims=3)) ==
+           @inferred(dropdims(sum, a, dims=(3,))) ==
+           reshape(sum(a, dims=3), (1, 1, 8, 1)))
+    @test (@inferred(dropdims(sum, a, dims=4)) ==
+           @inferred(dropdims(sum, a, dims=(4,))) ==
+           reshape(sum(a, dims=4), (1, 1, 8, 1)))
+    @test (@inferred(dropdims(sum, a, dims=(1, 5))) ==
+           dropdims(sum, a, dims=(5, 1)) ==
+           reshape(sum(a, dims=(5, 1)), (1, 8, 8)))
+    @test (@inferred(dropdims(sum, a, dims=(1, 2, 5))) ==
+           dropdims(sum, a, dims=(5, 2, 1)) ==
+           reshape(sum(a, dims=(5, 2, 1)), (8, 8)))
+    @test (@inferred(dropdims(sum, abs2, a, dims=1)) ==
+           @inferred(dropdims(sum, abs2, a, dims=(1,))) ==
+           reshape(sum(abs2, a, dims=1), (1, 8, 8, 1)))
+    _sumplus(x; dims, plus) = sum(x; dims=dims) .+ plus  # reduction with keywords
+    @test (@inferred(dropdims(_sumplus, a, dims=4, plus=1)) ==
+           @inferred(dropdims(_sumplus, a, dims=(4,), plus=1)) ==
+           reshape(sum(a, dims=4) .+ 1, (1, 1, 8, 1)))
+    @test_throws UndefKeywordError dropdims(sum, a)
+    @test_throws UndefKeywordError dropdims(sum, a, 1)
+    @test_throws ArgumentError dropdims(sum, a, dims=0)
+    @test_throws ArgumentError dropdims(sum, a, dims=(1, 1))
+    @test_throws ArgumentError dropdims(sum, a, dims=(1, 2, 1))
+    @test_throws ArgumentError dropdims(sum, a, dims=(1, 1, 2))
+    @test_throws ArgumentError dropdims(sum, a, dims=6)
+
     sz = (5,8,7)
     A = reshape(1:prod(sz),sz...)
     @test A[2:6] == [2:6;]