Understending `dims` kwarg

When I use mean function with dims=1 (the row dimension) in a matrix, the mean of each column is returned.
Example:

julia> using Statistics

julia> a = [3  5  4  5
            9  8  4  9
            9  9  5  6]
3×4 Matrix{Int64}:
 3  5  4  5
 9  8  4  9
 9  9  5  6

julia> mean(a, dims=1)
1×4 Matrix{Float64}:
 7.0  7.33333  4.33333  6.66667

Why does it happen? How does dims kwarg work?

The dim keyword argument specify in which dimension the n-dimensional array is going to be reduced. dims=1 means that the mean is going to be taken reducing over rows. dims=2 means reducing over columns.

For a 3D array, you can see how it works with this example

julia> a = [3;  5;;  4;  5;;;
            9;  8;;  4;  9;;;
            9;  9;;  5;  6]
2×2×3 Array{Int64, 3}:
[:, :, 1] =
 3  4
 5  5

[:, :, 2] =
 9  4
 8  9

[:, :, 3] =
 9  5
 9  6

julia> size(a) # rows, columns, depth
(2, 2, 3)

julia> mean(a; dims=1) # reduced over rows
1×2×3 Array{Float64, 3}:
[:, :, 1] =
 4.0  4.5

[:, :, 2] =
 8.5  6.5

[:, :, 3] =
 9.0  5.5

julia> mean(a; dims=2) # reduced over columns
2×1×3 Array{Float64, 3}:
[:, :, 1] =
 3.5
 5.0

[:, :, 2] =
 6.5
 8.5

[:, :, 3] =
 7.0
 7.5

julia> mean(a; dims=3) # reduced over depth
2×2×1 Array{Float64, 3}:
[:, :, 1] =
 7.0      4.33333
 7.33333  6.66667

Hope it makes it clear. :sweat_smile:

4 Likes