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

NEWS.md

@@ -30,16 +30,18 @@ New library functions
 * 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]).
 * The new `only(x)` function returns the one-and-only element of a collection `x`, and throws an `ArgumentError` if `x` contains zero or multiple elements. ([#33129])
 
-
 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]).
 
 * `div` now accepts a rounding mode as the third argument, consistent with the corresponding argument to `rem`. Support for rounding division, by passing one of the RoundNearest modes to this function, was added. For future compatibility, library authors should now extend this function, rather than extending the two-argument `fld`/`cld`/`div` directly. ([#33040])
 
+* 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]).
+
 * Verbose `display` of `Char` (`text/plain` output) now shows the codepoint value in standard-conforming `"U+XXXX"` format ([#33291]).
 
+
 #### 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;]