OpInfo for nn.functional.softmax #62077

krshrimali · 2021-07-23T05:25:25Z

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

facebook-github-bot · 2021-07-23T05:25:31Z

🔗 Helpful links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/62077
📄 Preview docs built from this PR

💊 CI failures summary and remediations

As of commit bd9a1d3 (more details on the Dr. CI page):

💚 💚 Looks good so far! There are no failures yet. 💚 💚

This comment was automatically generated by Dr. CI (expand for details).

Follow this link to opt-out of these comments for your Pull Requests.

Please report bugs/suggestions to the (internal) Dr. CI Users group.

Click here to manually regenerate this comment.

…nal/softmax

torch/testing/_internal/common_methods_invocations.py

mruberry · 2021-07-24T10:20:49Z

This looks pretty good @krshrimali; I have one suggestion (inline) for how to tweak testing the "dtype" kwarg.

krshrimali

A few updates:

This PR now uses sample inputs function of log_softmax for softmax.
Added alias for softmax.
Code clean-up for some OpInfos, when params are passed with default values (which isn't needed).
Skip removal for test_jit_alias_remapping test of log_softmax.

cc: @mruberry @zou3519 (sorry for the ping over the weekend, please review whenever you find time)

torch/testing/_internal/common_methods_invocations.py

krshrimali · 2021-07-26T06:03:42Z

Update:

I'm taking a look at the XLA error, if it's something non-trivial - I'll probably remove the scalar input and add an issue related to this error. Also, this is only reproducible on XLA (tested on Google Colab) with scalar tensors:

# Works on CPU
>>> torch.log_softmax(torch.rand((), device='cpu'), dim=0)
tensor(0.)

# Fails on XLA
>>> torch.log_softmax(torch.rand((), device='xla'), dim=0)
ERROR (please see the error below since it's too verbose)

Error:

Error on XLA

RuntimeError                              Traceback (most recent call last)
<ipython-input-4-470607feefbc> in <module>()
----> 1 torch.log_softmax(torch.rand((), device='xla'), dim=0)

RuntimeError: torch_xla/csrc/helpers.cpp:97 : Check failed: min_shape_dim <= dim && dim <= max_shape_dim 
*** Begin stack trace ***
	tensorflow::CurrentStackTrace()
	torch_xla::XlaHelpers::GetCanonicalDimensionIndex(long, long)
	torch_xla::XLATensor::log_softmax(torch_xla::XLATensor const&, long, c10::optional<c10::ScalarType>)
	torch_xla::AtenXlaType::_log_softmax(at::Tensor const&, long, bool)
	c10::impl::wrap_kernel_functor_unboxed_<c10::impl::detail::WrapFunctionIntoRuntimeFunctor_<at::Tensor (*)(at::Tensor const&, long, bool), at::Tensor, c10::guts::typelist::typelist<at::Tensor const&, long, bool> >, at::Tensor (at::Tensor const&, long, bool)>::call(c10::OperatorKernel*, c10::DispatchKeySet, at::Tensor const&, long, bool)
	at::Tensor c10::Dispatcher::redispatch<at::Tensor, at::Tensor const&, long, bool>(c10::TypedOperatorHandle<at::Tensor (at::Tensor const&, long, bool)> const&, c10::DispatchKeySet, at::Tensor const&, long, bool) const
	at::redispatch::_log_softmax(c10::DispatchKeySet, at::Tensor const&, long, bool)
	
	
	at::_log_softmax(at::Tensor const&, long, bool)
	at::native::log_softmax(at::Tensor const&, long, c10::optional<c10::ScalarType>)
	
	at::Tensor::log_softmax(long, c10::optional<c10::ScalarType>) const
	
	_PyMethodDef_RawFastCallKeywords
	_PyCFunction_FastCallKeywords
	_PyEval_EvalFrameDefault
	_PyEval_EvalCodeWithName
	PyEval_EvalCode
	
	_PyMethodDef_RawFastCallKeywords
	_PyCFunction_FastCallKeywords
	_PyEval_EvalFrameDefault
	_PyEval_EvalCodeWithName
	_PyFunction_FastCallKeywords
	_PyEval_EvalFrameDefault
	_PyEval_EvalCodeWithName
	_PyFunction_FastCallKeywords
	_PyEval_EvalFrameDefault
	_PyEval_EvalCodeWithName
	_PyObject_Call_Prepend
	PyObject_Call
	_PyEval_EvalFrameDefault
	_PyEval_EvalCodeWithName
	_PyFunction_FastCallKeywords
	_PyEval_EvalFrameDefault
	_PyEval_EvalCodeWithName
	_PyFunction_FastCallKeywords
	_PyEval_EvalFrameDefault
	_PyFunction_FastCallKeywords
	_PyEval_EvalFrameDefault
	_PyFunction_FastCallKeywords
	_PyEval_EvalFrameDefault
	_PyEval_EvalCodeWithName
	_PyFunction_FastCallDict
	_PyEval_EvalFrameDefault
	_PyEval_EvalCodeWithName
	_PyFunction_FastCallDict
	_PyEval_EvalFrameDefault
	_PyEval_EvalCodeWithName
	_PyFunction_FastCallKeywords
	_PyEval_EvalFrameDefault
	_PyFunction_FastCallKeywords
	_PyEval_EvalFrameDefault
	_PyObject_Call_Prepend
	PyObject_Call
	_PyEval_EvalFrameDefault
	_PyEval_EvalCodeWithName
	_PyFunction_FastCallKeywords
	_PyEval_EvalFrameDefault
	_PyObject_Call_Prepend
	_PyObject_FastCallKeywords
	
	_PyMethodDef_RawFastCallDict
	PyCFunction_Call
	_PyEval_EvalFrameDefault
	_PyFunction_FastCallKeywords
	_PyEval_EvalFrameDefault
	_PyFunction_FastCallKeywords
	_PyEval_EvalFrameDefault
	_PyFunction_FastCallKeywords
	_PyEval_EvalFrameDefault
	_PyFunction_FastCallKeywords
	_PyEval_EvalFrameDefault
	_PyFunction_FastCallKeywords
	_PyEval_EvalFrameDefault
	_PyEval_EvalCodeWithName
	_PyFunction_FastCallKeywords
	_PyEval_EvalFrameDefault
	_PyEval_EvalCodeWithName
	PyEval_EvalCode
	
	_PyMethodDef_RawFastCallKeywords
	_PyCFunction_FastCallKeywords
	_PyEval_EvalFrameDefault
	_PyEval_EvalCodeWithName
	_PyFunction_FastCallKeywords
	_PyEval_EvalFrameDefault
	_PyEval_EvalCodeWithName
	_PyFunction_FastCallDict
	
	
	_Py_UnixMain
	__libc_start_main
	_start
*** End stack trace ***
Value out of range (expected to be in range of [0, -1], but got 0)

Also, the error doesn't look message correct: since 0 is in the range [0, -1].

Will update once I've finished building XLA locally to test the macros used. Thanks!

krshrimali · 2021-07-27T14:42:06Z

Updates:

PyTorch on 0d tensors with dim=0 doesn't throw an error but on XLA, it does. This needs some discussion, but didn't want to block this PR hence I removed the input ((), 0) case (input shape, dim).

This PR should be ready for review, hopefully, tests should pass. Thanks!

cc: @mruberry @zou3519 @pmeier

pmeier

Two nits inline, otherwise LGTM! Thanks @krshrimali.

torch/testing/_internal/common_methods_invocations.py

Co-authored-by: Philip Meier <github.pmeier@posteo.de>

…nal/softmax

krshrimali · 2021-07-28T13:26:41Z

Gentle ping @zou3519 @mruberry - if you can take a look whenever you find the time. Thanks! :)

zou3519 · 2021-07-28T13:28:33Z

torch/testing/_internal/common_methods_invocations.py

-    def generator():
-        for shape, args, kwargs in cases:
-            yield SampleInput(make_arg(shape), args=args, kwargs=kwargs)
+    cases = (


This is a pathological case, but could we add ((), (0,)) to the list?

Thanks for the comment @zou3519! I had added this before but the test fails on the XLA device: #62077 (comment). PyTorch on 0d tensors with dim=0 doesn't throw an error but on XLA, it does.

Gotcha, sorry for not seeing that! The action items you proposed (file an issue, leave that case out of the OpInfo) sgtm

Thanks, @zou3519 ! Post discussion with @mruberry, we thought that having a separate case when the device type isn't XLA so that we don't skip this input for CPU/CUDA devices. I've also filed an issue here: pytorch/xla#3061.

torch/testing/_internal/common_methods_invocations.py

zou3519

This looks pretty good. I added some comments about some more cases for completeness, after that we should be good to go

…hrimali/pytorch into opinfo/nn/functional/softmax

…ntries now" This reverts commit 8ac9c63.

…nal/softmax

zou3519 · 2021-07-29T15:20:12Z

Lint is failing: https://github.com/pytorch/pytorch/pull/62077/checks?check_run_id=3191463113

Run (! git --no-pager grep -In '[[:blank:]]$' -- . ':(exclude)**/contrib/**' ':(exclude)**.diff' ':(exclude)third_party' || (echo "The above lines have trailing spaces; please remove them"; false))
torch/testing/_internal/common_methods_invocations.py:7922:        sample_inputs_func=partial(sample_inputs_softmax_variant, with_dtype=True), 
The above lines have trailing spaces; please remove them

krshrimali · 2021-07-29T15:28:48Z

Lint is failing: https://github.com/pytorch/pytorch/pull/62077/checks?check_run_id=3191463113

Run (! git --no-pager grep -In '[[:blank:]]$' -- . ':(exclude)**/contrib/**' ':(exclude)**.diff' ':(exclude)third_party' || (echo "The above lines have trailing spaces; please remove them"; false))
torch/testing/_internal/common_methods_invocations.py:7922:        sample_inputs_func=partial(sample_inputs_softmax_variant, with_dtype=True), 
The above lines have trailing spaces; please remove them

Thanks, @zou3519 for the pointer, I've fixed it now. Hopefully, the tests should pass :)

facebook-github-bot · 2021-07-29T15:28:54Z

@zou3519 has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

mruberry · 2021-07-30T04:40:58Z

torch/testing/_internal/common_methods_invocations.py

-        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.


Nice comment

mruberry · 2021-07-30T04:41:19Z

torch/testing/_internal/common_methods_invocations.py

-    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


Style nit: put the for loop first for readability (doesn't have to be changed in this PR)

mruberry · 2021-07-30T04:41:57Z

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)


"with_dtype" should be kwarg-only

mruberry · 2021-07-30T04:42:23Z

torch/testing/_internal/common_methods_invocations.py

@@ -6372,6 +6380,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,


Nice comment

mruberry

Really nice skip removal. Overall looks good. I made a few comments inline to consider in future PRs, no changes needed for this one

facebook-github-bot · 2021-07-30T13:58:11Z

@zou3519 merged this pull request in 09d10c4.

OpInfo for nn.functional.softmax

e675c2d

facebook-github-bot added the cla signed label Jul 23, 2021

krshrimali requested a review from mruberry July 23, 2021 05:25

Merge remote-tracking branch 'upstream/master' into opinfo/nn/functio…

3447920

…nal/softmax

pytorchbot added the open source label Jul 23, 2021

Added softmax to works_list, why didn't the CI raise an error though? :/

2e0cee2

krshrimali added the module: testing Issues related to the torch.testing module (not tests) label Jul 23, 2021

jbschlosser added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Jul 23, 2021

remove entry from works_list

e451c1c

mruberry reviewed Jul 24, 2021

View reviewed changes

torch/testing/_internal/common_methods_invocations.py Outdated Show resolved Hide resolved

mruberry requested a review from zou3519 July 24, 2021 10:20

krshrimali mentioned this pull request Jul 24, 2021

OpInfo Tracker #54261

Closed

krshrimali added 2 commits July 24, 2021 16:41

Code cleanup + use sample inputs func of log_softmax

28beaec

Rename sample inputs of log_softmax to softmax variant

e3fc338

krshrimali commented Jul 24, 2021

View reviewed changes

PyTorch vs XLA are diff for 0d tensors with dim arg

9eff904

make nn.functional.softmax as alias, use softmax in OpInfo

25b9357

pmeier approved these changes Jul 28, 2021

View reviewed changes

torch/testing/_internal/common_methods_invocations.py Outdated Show resolved Hide resolved

torch/testing/_internal/common_methods_invocations.py Show resolved Hide resolved

krshrimali and others added 3 commits July 28, 2021 16:17

Update torch/testing/_internal/common_methods_invocations.py

63e4936

Co-authored-by: Philip Meier <github.pmeier@posteo.de>

Merge remote-tracking branch 'upstream/master' into opinfo/nn/functio…

7a62c8f

…nal/softmax

narrow down this PR to the title

6ca9b04

zou3519 reviewed Jul 28, 2021

View reviewed changes

torch/testing/_internal/common_methods_invocations.py Outdated Show resolved Hide resolved

zou3519 reviewed Jul 28, 2021

View reviewed changes

krshrimali mentioned this pull request Jul 29, 2021

PyTorch Op Fails with 0d tensor and dim arg pytorch/xla#3061

Closed

krshrimali added 6 commits July 29, 2021 06:38

Add negative dim + include issue and separate case when not XLA

eb10d5d

Merge branch 'opinfo/nn/functional/softmax' of https://github.com/krs…

256ce49

…hrimali/pytorch into opinfo/nn/functional/softmax

Add with_lastdim to the sample inputs func, multiple opinfo entries now

8ac9c63

Revert "Add with_lastdim to the sample inputs func, multiple opinfo e…

075fad0

…ntries now" This reverts commit 8ac9c63.

Add bfloat16 for dtypesIfCPU

e3c4770

Merge remote-tracking branch 'upstream/master' into opinfo/nn/functio…

0af8aeb

…nal/softmax

Remove trailing space

bd9a1d3

zou3519 approved these changes Jul 29, 2021

View reviewed changes

mruberry reviewed Jul 30, 2021

View reviewed changes

mruberry self-requested a review July 30, 2021 04:43

mruberry approved these changes Jul 30, 2021

View reviewed changes

facebook-github-bot closed this in 09d10c4 Jul 30, 2021

facebook-github-bot added the Merged label Jul 30, 2021

zou3519 mentioned this pull request Jul 30, 2021

Top 25 OpInfos for functorch pytorch/functorch#78

Closed

23 tasks

OpInfo for nn.functional.softmax #62077

OpInfo for nn.functional.softmax #62077

Uh oh!

Conversation

krshrimali commented Jul 23, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

facebook-github-bot commented Jul 23, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful links

💊 CI failures summary and remediations

Uh oh!

Uh oh!

mruberry commented Jul 24, 2021

Uh oh!

krshrimali left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

krshrimali commented Jul 26, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

krshrimali commented Jul 27, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pmeier left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

krshrimali commented Jul 28, 2021

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

zou3519 left a comment

Choose a reason for hiding this comment

Uh oh!

zou3519 commented Jul 29, 2021

Uh oh!

krshrimali commented Jul 29, 2021

Uh oh!

facebook-github-bot commented Jul 29, 2021

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mruberry left a comment

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot commented Jul 30, 2021

Uh oh!

Uh oh!

krshrimali commented Jul 23, 2021 •

edited

Loading

facebook-github-bot commented Jul 23, 2021 •

edited

Loading

krshrimali left a comment •

edited

Loading

krshrimali commented Jul 26, 2021 •

edited

Loading

krshrimali commented Jul 27, 2021 •

edited

Loading