What is the best way to obtain every slice sums in tensor in Julia?

Question

T is a given tensor. I need the slice sums of T in every direction. For example, for a 4-order tensor T, I wrote the following code.

function get_every_slice_sums(T)
   A = sum(T, dims=[  2,3,4])
   B = sum(T, dims=[1,  3,4])
   C = sum(T, dims=[1,2,  4])
   D = sum(T, dims=[1,2,3  ])
   return A,B,C,D
end

But I think the above computation is redundant. So I also wrote

function get_every_slice_sums_fast(T)
   X = sum(T, dims=[3,4])
   Y = sum(T, dims=[1,2])

   A = sum(X, dims=[2])
   B = sum(X, dims=[1])
   C = sum(Y, dims=[4])
   D = sum(Y, dims=[3])
   return A,B,C,D
end

using BenchmarkTools

T = rand(20,20,20,20)
@btime get_every_slice_sums(T)
# 308.018 μs
@btime get_every_slice_sums_fast(T)
# 141.314 μs

It works. The second one is better than the original one. What is the best efficient way to do that for general order tensors?

EDIT

I am also looking for speed optimal algorithm.

Answer 1

Here is a basis for other possible answers. The following works for any Array shape, and is economical in memory accesses, but in the case of slice sums, not speed optimal:

function get_every_slice_sum(a::AbstractArray{T,D}) where {T,D}
    S = size(a)
    newT = eltype(foldl(+, fill(one(T),2)))
    res = ntuple(j->(zeros(newT, S[j])),D)
    for I in eachindex(IndexCartesian(), a)
        v = a[I]
        for j in 1:D
            res[j][I[j]] += v
        end
    end
    return res
end

and just in case, for fiber sums:

function get_every_fiber_sum(a::AbstractArray{T,D}) where {T,D}
    S = size(a)
    newT = eltype(foldl(+,fill(one(T),2)))
    res = ntuple(j->(zeros(newT, ntuple(i->(i==j ? S[i] : 1), D))),D)
    for I in eachindex(IndexCartesian(), a)
        v = a[I]
        for j in 1:D
            res[j][I[j]] += v
        end
    end
    return res
end