Repository: pytorch/pytorch

Issue:
Tensor.sum() over multiple axes
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
```
1. Is it ok?
2. Where should it be better placed?


Discussion:
Looks very useful

- Numpy uses `keepdims`, is there a reason to use `keepdim` name for parameter?
- probably, same should be done to `torch.prod`
- does it make sense to change `tensor.sum`, `variable.sum` as well?
Right, we should rename keepdim to keepdims if we do this (and support it for both variables and tensors).  Also should do it for all functions with a keepdim parameter.
If we want to add reduction over multiple axis, I think we should permute the axis so that they are consecutive, make contiguous and view those axis as one single dimension, so that we can perform the reduction over the reduced dimension. That would be more efficient than perfoming several reductions.
@fmassa, it depends. .contiguous() for 3+ dimensional tensors is very slow now. 
@ngimel hum, that's a good point. Also, we will rarely perform reductions over more than 3 axes, so the number of extra calls is small.
I've created a draft. Currently it checks number of dimensions. But it seems that choosing only one of the two (permute vs multiple calls to `.sum()`) solutions is better idea.

As @arogozhnikov mentioned: do we need to add similar implementation of `torch.prod`?
~We don't need whole tensor to be contiguous.~
~If we perform 3D `tensor.sum((1, 2))`, can reduce `sum` calls like~
```python
if tensor.stride(2)*tensor.size(2) == tensor.stride(1):
    output = tensor.view(tensor.size(0), -1).sum(1)
else:
    output = tensor.sum(2).sum(1)
```

And I think `mean`, `median`, `mode`, `min`, `max`, `std`, `var`, (and `norm`?) should have same signature for consistency.

**Edit:** I didn't know `view` can be performed only on contiguous tensor...
`contiguous` call is needed as @fmassa mentioned.

While we can avoid `contiguous` call like below.
```python
order = np.argsort(tensor.stride())[::-1].tolist()
permuted = tensor.permute(*order)

if permuted.is_contiguous():
    axes = tuple(map(lambda x: order[x], axes))
    # contiguous version
else:
    # not contiguous version
```

Resolving patch:
diff --git a/torch/_torch_docs.py b/torch/_torch_docs.py
--- a/torch/_torch_docs.py
+++ b/torch/_torch_docs.py
@@ -4094,7 +4094,8 @@ def parse_kwargs(desc):
 .. function:: sum(input, dim, keepdim=False, out=None) -> Tensor
 
 Returns the sum of each row of the :attr:`input` tensor in the given
-dimension :attr:`dim`.
+dimension :attr:`dim`. If :attr::`dim` is a list of dimensions,
+reduce over all of them.
 
 If :attr:`keepdim` is ``True``, the output tensor is of the same size
 as :attr:`input` except in the dimension :attr:`dim` where it is of size 1.
@@ -4103,7 +4104,7 @@ def parse_kwargs(desc):
 
 Args:
     input (Tensor): the input tensor
-    dim (int): the dimension to reduce
+    dim (int or tuple of ints): the dimension or dimensions to reduce
     keepdim (bool): whether the output tensor has :attr:`dim` retained or not
     out (Tensor, optional): the output tensor
 
@@ -4117,6 +4118,9 @@ def parse_kwargs(desc):
             [ 0.3637, -0.9906, -0.4752, -1.5197]])
     >>> torch.sum(a, 1)
     tensor([-0.4598, -0.1381,  1.3708, -2.6217])
+    >>> b = torch.arange(4 * 5 * 6).view(4, 5, 6)
+    >>> torch.sum(b, (2, 1))
+    tensor([  435.,  1335.,  2235.,  3135.])
 """)
 
 add_docstr(torch.svd,


Test patch:
diff --git a/test/test_autograd.py b/test/test_autograd.py
--- a/test/test_autograd.py
+++ b/test/test_autograd.py
@@ -2621,6 +2621,8 @@ class dont_convert(tuple):
     ('sum', (), NO_ARGS, 'scalar'),
     ('sum', (), (0,), 'scalar_dim', [0]),
     ('sum', (), (0, True,), 'scalar_keepdim_dim', [0]),
+    ('sum', (S, S, S), ([1, 2],), 'multi_dim'),
+    ('sum', (S, S, S), ([1, 2], True,), 'multi_dim_keepdim'),
     ('prod', (S, S, S), NO_ARGS),
     ('prod', (S, S, S), (1,), 'dim', [0]),
     ('prod', (S, S, S), (1, True,), 'keepdim_dim', [0]),
diff --git a/test/test_torch.py b/test/test_torch.py
--- a/test/test_torch.py
+++ b/test/test_torch.py
@@ -1503,6 +1503,8 @@ def make_tensors(*shape):
         check_sum_dim(make_tensors(50, 50, 50), 0)
         check_sum_dim(make_tensors(50, 50, 50), 1)
         check_sum_dim(make_tensors(50, 50, 50), 2)
+        check_sum_dim(make_tensors(50, 50, 50), (1, 2))
+        check_sum_dim(make_tensors(50, 50, 50), (1, -1))
 
         def make_contiguous_slice(size, dtype):
             contig = make_contiguous((1, size), dtype)
@@ -1522,6 +1524,11 @@ def test_sum_out(self):
         res2 = torch.Tensor()
         torch.sum(x, 1, out=res2)
         self.assertEqual(res1, res2)
+        x = torch.rand(100, 100, 100)
+        res1 = x.sum(2).sum(1)
+        res2 = torch.Tensor()
+        torch.sum(x, (2, 1), out=res2)
+        self.assertEqual(res1, res2)
 
     # TODO: these tests only check if it's possible to pass a return value
     # it'd be good to expand them