Repository: pytorch/pytorch

Issue:
Gradient formula for prod can't handle 0 inputs
They sometimes give NaNs or infs.


Discussion:

Resolving patch:
diff --git a/torch/autograd/_functions/reduce.py b/torch/autograd/_functions/reduce.py
--- a/torch/autograd/_functions/reduce.py
+++ b/torch/autograd/_functions/reduce.py
@@ -46,13 +46,45 @@ def forward(self, input):
     def backward(self, grad_output):
         if self.dim is None:
             input, = self.saved_tensors
-            grad_input = grad_output.new(self.input_size).fill_(self.result)
-            return grad_input.div(input)
+            zero_idx = (input == 0).nonzero()
+            if zero_idx.dim() == 0:
+                return grad_output.mul(self.result).expand_as(input).div(input)
+            elif zero_idx.size(0) > 1:
+                return grad_output.new(self.input_size).zero_()
+            else:
+                grad_input = grad_output.new(self.input_size).zero_()
+                zero_idx = tuple(zero_idx[0].cpu())
+                input_copy = input.clone()
+                input_copy[zero_idx] = 1.
+                grad_input[zero_idx] = grad_output[0] * input_copy.prod()
+                return grad_input
         else:
             input, output = self.saved_tensors
-            repeats = [1 for _ in self.input_size]
-            repeats[self.dim] = self.input_size[self.dim]
-            return output.mul(grad_output).repeat(*repeats).div_(input)
+            zero_mask = input == 0
+            slice_zero_count = zero_mask.sum(self.dim)
+            total_zeros = slice_zero_count.sum()
+            grad_input = grad_output.mul(output).expand_as(input).div(input)
+            if total_zeros == 0:
+                return grad_input
+
+            some_zeros = slice_zero_count.gt(0).expand_as(grad_input)
+            grad_input[some_zeros] = 0
+
+            single_zero_idx = slice_zero_count.eq(1).nonzero()
+            for idx in single_zero_idx:
+                idx_tuple = tuple(idx.cpu())
+                input_idx_tuple = idx_tuple[:self.dim] + (slice(0, None),) + idx_tuple[self.dim + 1:]
+
+                # slice_mask and input_copy are 1D
+                slice_mask = zero_mask[input_idx_tuple]
+                input_copy = input[input_idx_tuple].clone()
+                zero_idx = slice_mask.nonzero()[0, 0]
+                input_copy[zero_idx] = 1.
+
+                grad_idx_tuple = idx_tuple[:self.dim] + (zero_idx,) + idx_tuple[self.dim + 1:]
+                grad_input[grad_idx_tuple] = grad_output[idx_tuple] * input_copy.prod()
+
+            return grad_input
 
 
 class Mean(_DimReduceFunction):
diff --git a/torch/autograd/_functions/tensor.py b/torch/autograd/_functions/tensor.py
--- a/torch/autograd/_functions/tensor.py
+++ b/torch/autograd/_functions/tensor.py
@@ -113,15 +113,17 @@ def __init__(self, sizes):
 
     def forward(self, i):
         result = i.expand(*self.sizes)
-        unsqueezed = (1,) * (len(self.sizes) - len(i.size()))
+        self.num_unsqueezed = len(self.sizes) - i.dim()
         self.expanded_dims = [dim for dim, (expanded, original)
-                              in enumerate(zip(self.sizes, unsqueezed + i.size()))
+                              in enumerate(zip(self.sizes[self.num_unsqueezed:], i.size()))
                               if expanded != original]
         self.mark_shared_storage((i, result))
         return result
 
     def backward(self, grad_output):
         grad_input = grad_output
+        for i in range(self.num_unsqueezed):
+            grad_input = grad_input.sum(0).squeeze(0)
         for dim in self.expanded_dims:
             grad_input = grad_input.sum(dim)
         return grad_input
diff --git a/torch/autograd/gradcheck.py b/torch/autograd/gradcheck.py
--- a/torch/autograd/gradcheck.py
+++ b/torch/autograd/gradcheck.py
@@ -146,4 +146,15 @@ def fn(input):
         for a, n in zip(analytical, numerical):
             if not ((a - n).abs() <= (atol + rtol * n.abs())).all():
                 return False
+
+    # check if the backward multiplies by grad_output
+    zero_gradients(inputs)
+    output = _as_tuple(func(*inputs))
+    torch.autograd.backward(output, [o.data.new(o.size()).zero_() for o in output])
+    for i in inputs:
+        if i.grad is None:
+            continue
+        if not i.grad.data.eq(0).all():
+            return False
+
     return True


Test patch:
diff --git a/test/test_autograd.py b/test/test_autograd.py
--- a/test/test_autograd.py
+++ b/test/test_autograd.py
@@ -940,6 +940,20 @@ def gather_variable(shape, index_dim, max_indices):
     return Variable(index, requires_grad=False)
 
 
+def prod_zeros(dim_size):
+    result = torch.randn(dim_size, dim_size, dim_size)
+    result[0, 1] = 0
+    result[2, 3] = 0
+    result[4, 3] = 0
+    return Variable(result, requires_grad=True)
+
+
+def prod_single_zero(dim_size):
+    result = torch.randn(dim_size, dim_size)
+    result[0, 1] = 0
+    return Variable(result, requires_grad=True)
+
+
 L = 20
 M = 10
 S = 5
@@ -1003,7 +1017,10 @@ def gather_variable(shape, index_dim, max_indices):
     (Sum, (), ((S, S, S),)),
     (Sum, (1,), ((S, S, S),), 'dim'),
     (Prod, (), ((S, S, S),)),
+    (Prod, (), (prod_zeros(S),), 'zeros'),
+    (Prod, (), (prod_single_zero(S),), 'single_zero'),
     (Prod, (1,), ((S, S, S),), 'dim'),
+    (Prod, (1,), (prod_zeros(S),), 'zeros_dim'),
     (Addmm, (), ((S, M), (S, S), (S, M)),),
     (Addmm, (0.1, 1), ((S, M), (S, S), (S, M)), 'coef'),
     (Addbmm, (), ((S, M), (S, S, S), (S, S, M)),),