Implement LogSoftmax double backwards (pytorch#2270)

Author: gchanan

test/common_nn.py

@@ -101,14 +101,12 @@
         module_name='LogSoftmax',
         input_size=(10, 20),
         reference_fn=lambda i, _: torch.exp(i).div_(torch.exp(i).sum(1, True).expand(10, 20)).log_(),
-        check_gradgrad=False,
     ),
     dict(
         module_name='LogSoftmax',
         input_size=(1, 3, 10, 20),
         reference_fn=lambda i, _: torch.exp(i).div_(torch.exp(i).sum(1, False)).log_(),
         desc='multiparam',
-        check_gradgrad=False,
     ),
     dict(
         module_name='ELU',

torch/nn/_functions/thnn/auto_double_backwards.py

@@ -39,6 +39,21 @@ def leakyrelu_double_backwards(ctx, ggI):
     return gI, ggO, None, None, None
 
 
+def logsoftmax_double_backwards(ctx, ggI):
+    t = ctx.saved_variables
+    gO, output = t[1], t[2]
+
+    output_exp = output.exp()
+    gO_sum = gO.sum(dim=1, keepdim=True)
+    ggI_output_exp = ggI * output_exp
+    ggI_output_exp_sum = ggI_output_exp.sum(dim=1, keepdim=True)
+
+    gI = output_exp * gO_sum * ggI_output_exp_sum - ggI_output_exp * gO_sum
+    ggO = ggI - ggI_output_exp_sum
+
+    return gI, ggO, None, None, None, None
+
+
 def softmax_double_backwards(ctx, ggI):
     t = ctx.saved_variables
     gO, output = t[1], t[2]
@@ -126,6 +141,7 @@ def nllloss_double_backwards(ctx, ggI):
     'ELU': elu_double_backwards,
     'Hardtanh': hardtanh_double_backwards,
     'LeakyReLU': leakyrelu_double_backwards,
+    'LogSoftmax': logsoftmax_double_backwards,
     'Softmax': softmax_double_backwards,
     'Threshold': threshold_double_backwards,
     'L1Loss': l1loss_double_backwards,