在 Julia 中获取张量中每个切片总和的最佳方法是什么？

Question

T

是给定的张量。我需要每个方向上

T

的切片和。例如，对于4阶张量

T

，我编写了以下代码。

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

但是我认为上面的计算是多余的。所以我也写了

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

它有效。第二个比第一个好。对于一般阶张量，最有效的方法是什么？

Answer 1

这是其他可能答案的基础。以下适用于任何数组形状，并且在内存访问方面很经济，但在切片和的情况下，速度不是最佳的：

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

为了以防万一，对于纤维总量：

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