Tensor.sum() over multiple axes #2006
Description
During writing einsum I've found that sum is able to sum only over single axis. I'm going to add a simple implementation of it:
def sum(input, axes, keepdim=False):
# probably some check for uniqueness of axes
if keepdim:
for ax in axes:
input = input.sum(ax, keepdim=True)
else:
for ax in sorted(axes, reverse=True):
input = input.sum(ax)
return input
- Is it ok?
- Where should it be better placed?