OpInfo for nn.functional.softmax (#62077)

krshrimali · facebook-github-bot · commit 09d10c43296b · 2021-07-30T06:56:03.000-07:00
Summary: This PR: * Adds OpInfo for `softmax` and `nn.functional.softmax` (alias). * Skip removal for `test_jit_alias_remapping` test of `log_softmax`. Please see pytorch/functorch#78 and #54261. cc: mruberry zou3519 pmeier Pull Request resolved: #62077 Reviewed By: heitorschueroff Differential Revision: D29990019 Pulled By: zou3519 fbshipit-source-id: 67476990b54a5dd824eed9d10236e118564f2501
diff --git a/torch/testing/_internal/common_methods_invocations.py b/torch/testing/_internal/common_methods_invocations.py
@@ -3669,19 +3669,27 @@ def sample_inputs_to_sparse(op_info, device, dtype, requires_grad, **kwargs):
             SampleInput(make_arg((S, S)), args=(1,), output_process_fn_grad=lambda x: x.to_dense()),)
 
 
-def sample_inputs_log_softmax(op_info, device, dtype, requires_grad, with_dtype=False, **kwargs):
+# Used for both log_softmax and softmax
+def sample_inputs_softmax_variant(op_info, device, dtype, requires_grad, with_dtype=False, **kwargs):
     make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad)
 
-    if with_dtype:
-        cases = (((S, S, S), (1, ), {'dtype': torch.float64}),)
-    else:
-        cases = (((S, S, S), (1,), {}),)  # type:ignore[assignment]
+    cases = [
+        ((S, ), (0, )),
+        ((S, S), (0, )),
+        ((S, S), (1, )),
+        ((S, S), (-1, )),
+        ((S, M, S), (2, )),
+    ]
 
-    def generator():
-        for shape, args, kwargs in cases:
-            yield SampleInput(make_arg(shape), args=args, kwargs=kwargs)
+    # PyTorch on XLA throws an error when passed with dim argument for 0d tensor.
+    # See https://github.com/pytorch/xla/issues/3061 for more details.
+    if torch.device(device).type != 'xla':
+        cases.append(((), (0, )))
 
-    return list(generator())
+    return [
+        SampleInput(make_arg(shape), args=dim, kwargs=dict(dtype=torch.float64) if with_dtype else None)
+        for shape, dim in cases
+    ]
 
 
 def sample_inputs_logit(op_info, device, dtype, requires_grad, **kwargs):
@@ -6374,6 +6382,24 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs):
            dtypes=all_types_and(torch.float16, torch.bfloat16, torch.bool),
            supports_forward_ad=True,
            sample_inputs_func=sample_inputs_max_min_binary,),
+    # `softmax` supports different dtypes based on whether `dtype` argument,
+    # is passed or not. Hence two OpInfo entries, one with dtype and other without.
+    OpInfo('softmax',
+           aliases=('nn.functional.softmax',),
+           aten_name='softmax',
+           dtypesIfCPU=floating_types_and(torch.bfloat16),
+           dtypesIfCUDA=floating_types_and(torch.half, torch.bfloat16),
+           sample_inputs_func=sample_inputs_softmax_variant,
+           assert_autodiffed=True,
+           supports_out=False),
+    OpInfo('softmax',
+           aliases=('nn.functional.softmax',),
+           variant_test_name="with_dtype",
+           aten_name='softmax',
+           dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16),
+           sample_inputs_func=partial(sample_inputs_softmax_variant, with_dtype=True),
+           assert_autodiffed=True,
+           supports_out=False),
     OpInfo('nn.functional.hardswish',
            aten_name="hardswish",
            supports_autograd=True,
@@ -7887,32 +7913,15 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs):
         supports_out=False,
         dtypes=floating_types_and(torch.bfloat16),
         dtypesIfCUDA=floating_types_and(torch.float16, torch.bfloat16),
-        sample_inputs_func=sample_inputs_log_softmax,
+        sample_inputs_func=sample_inputs_softmax_variant,
         assert_autodiffed=True),
     OpInfo(
         'log_softmax',
         variant_test_name='dtype',
         aliases=('special.log_softmax', 'nn.functional.log_softmax'),
         supports_out=False,
         dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16),
-        sample_inputs_func=partial(sample_inputs_log_softmax, with_dtype=True),
-        skips=(
-            # NOTE: This should work once https://github.com/pytorch/pytorch/pull/58838 is in
-            # RuntimeError:
-            # Unknown type name 'dtype':
-            # File "<string>", line 2
-            #         def _fn(t0, s0: int, dtype: dtype = torch.float64):
-            #                                     ~~~~~ <--- HERE
-            #             return variant(t0, s0, dtype=torch.float64)
-
-            # 'defaults' is being compiled since it was called from '_fn'
-            # File "<string>", line 2
-            #         def _fn(t0, s0: int, dtype: dtype = torch.float64):
-            #             ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-            #             return variant(t0, s0, dtype=torch.float64)
-            #             ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
-            SkipInfo('TestJit', 'test_jit_alias_remapping'),
-        ),
+        sample_inputs_func=partial(sample_inputs_softmax_variant, with_dtype=True),
         assert_autodiffed=True),
     UnaryUfuncInfo('logit',
                    ref=scipy.special.logit if TEST_SCIPY else _NOTHING,
