如何计算ndarray上一维的平方和？

Question

一个多维矩阵形状（2,50,25,3）：

xx = np.random.randn(2, 50, 25, 3)

我想计算最后一个维度的平方和。结果应该是具有形状（2,50,25,1）的矩阵。

[np.sum(x) for x in np.square(features_displacement[0][0][:])[:]]

此代码可以成功计算一维，输出一个形状为（25,1）的列表，但如何计算如上所述的所有维度？

Answer 1

您可以沿所需的轴应用numpy函数，例如：

np.sum(np.square(xx), axis=3)

将产生一个形状(2, 50, 25)阵列。不完全确定这是你想要的，如果没有请更具体:-)

Answer 2

用这个：

sum_of_squares = np.sum(np.square(features_displacement), axis=-1)