《动手学深度学习》 05线性代数

按特定的 轴 axis 求和:

默认情况下，调用求和函数会沿所有的轴降低张量的维度，使它变为一个标量。 我们还可以[指定张量沿哪一个轴来通过求和降低维度]。 以矩阵为例，为了通过求和所有行的元素来降维（轴0），可以在调用函数时指定axis=0。 由于输入矩阵沿0轴降维以生成输出向量，因此输入轴0的维数在输出形状中消失。

A_sum_axis0 = A.sum(axis=0)
A_sum_axis0, A_sum_axis0.shape,A

(tensor([40., 45., 50., 55.]),
 torch.Size([4]),
 tensor([[ 0.,  1.,  2.,  3.],
         [ 4.,  5.,  6.,  7.],
         [ 8.,  9., 10., 11.],
         [12., 13., 14., 15.],
         [16., 17., 18., 19.]]))

指定axis=1将通过汇总所有列的元素降维（轴1）。因此，输入轴1的维数在输出形状中消失。

A_sum_axis1 = A.sum(axis=1)
A_sum_axis1, A_sum_axis1.shape

(tensor([ 6., 22., 38., 54., 70.]), torch.Size([5]))

• 可以用 keepdims 保留 sum原来的维度

A = torch.ones([2,4,5])
A.shape
#torch.Size([2, 4, 5])
 
A.sum(axis=(0,2))
#tensor([10., 10., 10., 10.])

A.sum(axis=(0,2),keepdims=True)
#tensor([[[10.],
         [10.],
         [10.],
         [10.]]])
         
A.sum(axis=(0,2),keepdims=True).shape         
#torch.Size([1, 4, 1])