Add autograd tests for keepdim

Author: soumith

test/test_autograd.py

@@ -1254,8 +1254,10 @@ class dont_convert(tuple):
     (CminConstant, (), ((S, S, S), 0.5)),
     (Mean, (), ((S, S, S),)),
     (Mean, (1,), ((S, S, S),), 'dim', [0]),
+    (Mean, (1, True,), ((S, S, S),), 'keepdim_dim', [0]),
     (Sum, (), ((S, S, S),)),
     (Sum, (1,), ((S, S, S),), 'dim', [0]),
+    (Sum, (1, True,), ((S, S, S),), 'keepdim_dim', [0]),
     (Prod, (), ((S, S, S),)),
     (Prod, (), (prod_zeros(S, [0, 1]),), 'zerosdim2'),
     (Prod, (), (prod_zeros(S, [0, 2]),), 'zerosdim1'),
@@ -1265,6 +1267,10 @@ class dont_convert(tuple):
     (Prod, (1,), (prod_zeros(S, [0, 1]),), 'zeros_dim2', [0]),
     (Prod, (1,), (prod_zeros(S, [0, 2]),), 'zeros_dim1', [0]),
     (Prod, (1,), (prod_zeros(S, [1, 2]),), 'zeros_dim0', [0]),
+    (Prod, (1, True,), ((S, S, S),), 'keepdim_dim', [0]),
+    (Prod, (1, True,), (prod_zeros(S, [0, 1]),), 'keepdim_zeros_dim2', [0]),
+    (Prod, (1, True,), (prod_zeros(S, [0, 2]),), 'keepdim_zeros_dim1', [0]),
+    (Prod, (1, True), (prod_zeros(S, [1, 2]),), 'keepdim_zeros_dim0', [0]),
     (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)),),
@@ -1284,15 +1290,23 @@ class dont_convert(tuple):
     (Min, (), ((S, S, S),),),
     (Max, (1,), ((S, S, S),), 'dim', [0]),
     (Min, (1,), ((S, S, S),), 'dim', [0]),
+    (Max, (1, True), ((S, S, S),), 'keepdim_dim', [0]),
+    (Min, (1, True), ((S, S, S),), 'keepdim_dim', [0]),
     (Mode, (1,), ((S, S, S),), 'dim', [0]),
+    (Mode, (1, True,), ((S, S, S),), 'keepdim_dim', [0]),
     (Kthvalue, (2, 0), ((S, S, S),),),
+    (Kthvalue, (2, 0, True), ((S, S, S),), "keepdim"),
     (Median, (0,), ((S, S, S),),),
+    (Median, (0, True, ), ((S, S, S),), "keepdim"),
     (Norm, (1.5,), (torch.rand(S, S, S),), '1_5'),
     (Norm, (), ((S, S, S),), '2'),
     (Norm, (3,), ((S, S, S),), '3'),
     (Norm, (1.5, 1), (torch.rand(S, S, S),), '1_5_dim', [1]),
     (Norm, (2, 1), ((S, S, S),), '2_dim', [1]),
     (Norm, (3, 1), ((S, S, S),), '3_dim', [1]),
+    (Norm, (1.5, 1, True,), (torch.rand(S, S, S),), 'keepdim_1_5_dim', [1]),
+    (Norm, (2, 1, True,), ((S, S, S),), 'keepdim_2_dim', [1]),
+    (Norm, (3, 1, True), ((S, S, S),), 'keepdim_3_dim', [1]),
     (Addcmul, (), ((S, S), (S, S), (S, S))),
     (Addcmul, (0.6,), ((S, S), (S, S), (S, S)), 'scale'),
     (Addcdiv, (), ((S, S), (S, S), torch.rand(S, S) + 5e-2)),
@@ -1388,20 +1402,31 @@ class dont_convert(tuple):
     ('lerp', (S, S, S), ((S, S, S), 0.4)),
     ('max', (S, S, S), ()),
     ('max', (S, S, S), (1,), 'dim', [0]),
+    ('max', (S, S, S), (1, True,), 'keepdim_dim', [0]),
     ('max', (S, S, S), ((S, S, S),), 'elementwise'),
     ('min', (S, S, S), ()),
     ('min', (S, S, S), (1,), 'dim', [0]),
+    ('min', (S, S, S), (1, True,), 'keepdim_dim', [0]),
     ('min', (S, S, S), ((S, S, S),), 'elementwise'),
     ('mean', (S, S, S), ()),
     ('mean', (S, S, S), (1,), 'dim', [0]),
+    ('mean', (S, S, S), (1, True,), 'keepdim_dim', [0]),
+    ('median', (S, S, S), (1,), 'dim', [0]),
+    ('median', (S, S, S), (1, True,), 'keepdim_dim', [0]),
+    ('mode', (S, S, S), (1,), 'dim', [0]),
+    ('mode', (S, S, S), (1, True,), 'keepdim_dim', [0]),
     ('sum', (S, S, S), ()),
     ('sum', (S, S, S), (1,), 'dim', [0]),
+    ('sum', (S, S, S), (1, True,), 'keepdim_dim', [0]),
     ('prod', (S, S, S), ()),
     ('prod', (S, S, S), (1,), 'dim', [0]),
+    ('prod', (S, S, S), (1, True,), 'keepdim_dim', [0]),
     ('var', (S, S, S), ()),
     ('var', (S, S, S), (1,), 'dim', [0]),
+    ('var', (S, S, S), (1, True), 'keepdim_dim', [0]),
     ('std', (S, S, S), ()),
     ('std', (S, S, S), (1,), 'dim', [0]),
+    ('std', (S, S, S), (1, True), 'keepdim_dim', [0]),
     ('renorm', (S, S, S), (2, 1, 0.5), 'dim', [1]),
     ('renorm', (S, S, S), (1, 2, 3), 'norm_1'),
     ('repeat', (S, S, S, S), (2, 3, 1, 4)),
@@ -1424,6 +1449,7 @@ class dont_convert(tuple):
     ('addcdiv', (S, S), (0.5, (S, S), (S, S)), 'scale'),
     ('norm', (S, S, S), (2,)),
     ('norm', (S, S, S), (2, 1), 'dim', [1]),
+    ('norm', (S, S, S), (2, 1, True), 'keepdim_dim', [0]),
     ('dist', (S, S, S), ((S, S, S),)),
     ('dist', (S, S, S), ((S, S, S), 4), '4'),
     ('index_select', (S, S, S), (0, index_variable(2, S)), 'dim', [0]),