How to calculate the sums of squares along one dimension of on ndarray?

Question

One miltidimensional matrix with shape (2, 50, 25, 3):

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

I want to calculate the sum of squares of the last dimension. The result should be a matrix with a shape (2, 50, 25, 1).

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

This code can successfully calculate the one dimension, output a list with shape (25,1), but how can I calculate all the dimensions as described above?

Answer 1

You can apply the numpy functions along the axis you want, for example:

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

Will produce an array of shape (2, 50, 25) . Not exactly sure this is what you want, if not please be more specific :-)

Answer 2

用这个：

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