Specialize optional tensor inputs to graphs in the JIT #18360

t-vi · 2019-03-22T20:39:22Z

This specializes optional tensor inputs to either a DimensionedTensorType or, when None is passed,
UndefinedTensor (aka AutogradZeroTensorType).
This works because we already have different specs and thus separate plans for the two cases.
It enhances the shape analysis - because now unwrapped optional tensors will have DimensionedTensorType with appropriate shape and required grad etc.
Also, when combined with "if-pruning" (which I understand #18259 works towards), we actually get much nicer concrete graphs, too.

This specializes optional tensor inputs to either a DimensionedTensorType or, when None is passed, UndefinedTensor (aka AutogradZeroTensorType). This works because we already have different specs and thus separate plans for the two cases. It enhances the shape analysis - because now unwrapped optional tensors will have DimensionedTensorType with appropriate shape and required grad etc. Also, when combined with "if-pruning" (which I understand

eellison · 2019-03-22T22:13:28Z

Looks great!

I'm going to wait for @wanchaol to review since I have less of a knowledge of how UndefinedTensor works. I think maybe mustBeNone() api should check for UndefinedTensor case now but I'm not completely sure

t-vi · 2019-03-22T22:14:35Z

Well, I'll follow up if our patches don't achieve what I have in mind. :)

eellison · 2019-03-22T22:15:19Z

Also i'm not completely sure on this either but Constant Propagation / Constant Pooling might need to be updated now to account for this case now too

t-vi · 2019-03-22T22:16:36Z

Adding the input shapes happens right at the beginning, so I think it would be before the various passes.

eellison

I think this looks good. I will update my other PR so that it handles the Undefined Tensor case. Maybe wait to land until @wanchaol comments

eellison · 2019-03-22T22:36:45Z

test/test_jit.py

+                return 0
+
+            fn(None)
+            g = fn.graph_for(None)


Add a check that the output of the function is 1 / 0 ?

Well, I only added a return value to not have it tell me that it needs to do nothing... :)

ya more just as a sanity check that the optimizations i'm adding in the other PR don't break this, although i'll wait for the landing of that PR with this one to be safe

What I would suggest is that we amend the check to see that the graph is only a constant return for each case. :)

wanchaol

lgtm, have a quick conversation with @eellison, this might need to deal with other types for shape analysis.

eellison · 2019-03-22T22:52:51Z

torch/csrc/jit/passes/shape_analysis.cpp

@@ -513,6 +513,12 @@ class ShapePropagator {
        }
        return;
      }
+      case prim::unchecked_unwrap_optional: {
+	if (node->input()->type()->isSubtypeOf(TensorType::get())) {


The more general case here is to set the output to be the unwrapped type of the input if it's an optional input, and the input if it's not an optional

I think we need to be a bit more detailed than that about None. So I'd use the following:

if it's an IValue and that isNone, leave the output alone,

if it's an OptionalType return the elementType,

otherwise return the Type.

Well, for unchecked, we could always pass on the type because it should only be inserted where we know it's not None, so I'll skip the first there.

eellison · 2019-03-22T22:53:54Z

torch/csrc/jit/passes/shape_analysis.cpp

@@ -529,6 +535,10 @@ class ShapePropagator {
        return;
      }
      case aten::_unwrap_optional: {
+	if (node->input()->type()->isSubtypeOf(TensorType::get())) {


Since _unwrap_optional will throw an exception if it's None, we can do the same thing as with prim::unchecked_unwrap and set the output to be the unwrapped type of the input if it's an optional input, and the input if it's not an optional

Same as above. Here, we need the None handling and can use the form that is already in that case.

Conditions as discussed in the review With correct indenting, the test actually tests...

t-vi · 2019-03-23T12:57:16Z

The CI failures seem unrelated, so if you're happy with it, I'd suggest to land this.

facebook-github-bot

@soumith is landing this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

t-vi · 2019-03-23T18:32:19Z

So Adam thinks it that it is preferrable to use None instead of UndefinedTensor. I'll follow up with an improvement.

In pytorch#18360, we used undefined Tensor (aka AutogradZeroTensor), but this can be errorprone when the type or value is compared to None, e.g. as seen when comined with the (not yet landed) For this to work, we must allow None passed to functions taking Tensor?.

apaszke

I'm very confused by what has changed here

apaszke · 2019-03-24T19:25:45Z

torch/csrc/jit/argument_spec.h

@@ -153,7 +153,8 @@ struct ArgumentSpec {

 private:
  TypePtr fillType(TypePtr original, size_t& offset) const {
-    if (original->isSubtypeOf(TensorType::get())) {
+    if (original->isSubtypeOf(TensorType::get())
+	|| original->isSubtypeOf(OptionalType::ofTensor())) {


This is quite surprising. Do we seriously translate None to an undefined tensor when you call a script function? How can that even work?

Oh, is it because all args are zero-initialized, so you're abusing the fact that this arg is not a tensor, but it will report being undefined? Bleh

apaszke · 2019-03-24T19:26:32Z

torch/csrc/jit/passes/shape_analysis.cpp

+	if(auto ot = node->input()->type()->cast<OptionalType>()) {
+	  node->output()->setType(ot->getElementType());
+	} else {
+	  node->output()->setType(node->input()->type());


Why would we have those calls on non-optional inputs?

Also, when is that op used?

These calls would be inserted before input types have been specialized. If we have an optional tensor as an input to the graph and we specialize the graph to a non-optional tensor, than these calls will no longer be needed

When does an unchecked_unwrap_optional ever have a different type than ot->getElementType? This code looks like a no-op to me. Why is it ineeded?

So as this PR is dead and has been reverted, I tried to add a better high-level overview this in #18407 .
The gist is that I've been trying to convert Optional[Tensor] into DimensionedTensorType or NoneType once we know what has been input.
May I suggest we move the discussion to #18407?

apaszke · 2019-03-24T19:28:09Z

torch/csrc/jit/passes/shape_analysis.cpp

+	  node->output()->setType(ot->getElementType());
+	} else {
+	  node->output()->setType(node->input()->type());
+	}


Why did we add this? How is _unwrap_optional different from unchecked_unwrap_optional? This code looks really weird, because we're trying to see if we can infer statically that this thing gets a None as an argument (which it never should??), and otherwise we assume it doesn't?

Ok. I think we should revert this and probably also abandon the newer patch.

_unwrap_optional is user code. it throws an error if the input is None and returns the input. unchecked_unwrap_optional has been inserted by the compiler when we can statically reason that the input will never be none (and so doesn't check if the input is None)

…ch#18360)" This reverts commit 7cc7ed1.

…" (#18411) Summary: This reverts commit 7cc7ed1. I think it's better to sort out the issues raised in #18407 firs. I'm sorry for not stopping it earlier. Pull Request resolved: #18411 Differential Revision: D14594937 Pulled By: soumith fbshipit-source-id: 3c90b7fa7694e2f59e55607acecde4a47af801ea

facebook-github-bot added the oncall: jit Add this issue/PR to JIT oncall triage queue label Mar 22, 2019

t-vi requested a review from eellison March 22, 2019 20:39

t-vi changed the title ~~Specialize optional tensor inputs to graphs in the JIT~~ [WIP] Specialize optional tensor inputs to graphs in the JIT Mar 22, 2019

don't fall through switch!

27d15cf

t-vi changed the title ~~[WIP] Specialize optional tensor inputs to graphs in the JIT~~ Specialize optional tensor inputs to graphs in the JIT Mar 22, 2019

t-vi requested a review from wanchaol March 22, 2019 22:15

eellison approved these changes Mar 22, 2019

View reviewed changes

wanchaol approved these changes Mar 22, 2019

View reviewed changes

eellison reviewed Mar 22, 2019

View reviewed changes

Refine conditions, fix test

5006fa4

Conditions as discussed in the review With correct indenting, the test actually tests...

facebook-github-bot reviewed Mar 23, 2019

View reviewed changes

facebook-github-bot closed this in 7cc7ed1 Mar 24, 2019

pytorchbot added the merged label Mar 24, 2019

t-vi mentioned this pull request Mar 24, 2019

[JIT] JIT erroneously sets requires grad #18270

Closed

t-vi mentioned this pull request Mar 24, 2019

Specialize Optional[T] to T (or subtype for Tensor) or None when executing graph #18407

Closed

apaszke reviewed Mar 24, 2019

View reviewed changes

t-vi added a commit to t-vi/pytorch that referenced this pull request Mar 24, 2019

Revert "Specialize optional tensor inputs to graphs in the JIT (pytor…

705a4f1

…ch#18360)" This reverts commit 7cc7ed1.

eellison mentioned this pull request Mar 25, 2019

Optimize boolean expressions & unwraps #18259

Closed

ezyang added the open source label Jun 24, 2019

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Specialize optional tensor inputs to graphs in the JIT #18360

Specialize optional tensor inputs to graphs in the JIT #18360

t-vi commented Mar 22, 2019

eellison commented Mar 22, 2019 •

edited

t-vi commented Mar 22, 2019

eellison commented Mar 22, 2019

t-vi commented Mar 22, 2019

eellison left a comment

eellison Mar 22, 2019

t-vi Mar 22, 2019

eellison Mar 22, 2019 •

edited

t-vi Mar 23, 2019

wanchaol left a comment

eellison Mar 22, 2019

t-vi Mar 23, 2019

t-vi Mar 23, 2019 •

edited

eellison Mar 22, 2019

t-vi Mar 23, 2019 •

edited

t-vi commented Mar 23, 2019

facebook-github-bot left a comment

t-vi commented Mar 23, 2019

apaszke left a comment

apaszke Mar 24, 2019

apaszke Mar 24, 2019

apaszke Mar 24, 2019

apaszke Mar 24, 2019

eellison Mar 25, 2019

zdevito Mar 25, 2019

t-vi Mar 25, 2019

apaszke Mar 24, 2019

t-vi Mar 24, 2019

eellison Mar 25, 2019

Specialize optional tensor inputs to graphs in the JIT #18360

Specialize optional tensor inputs to graphs in the JIT #18360

Conversation

t-vi commented Mar 22, 2019

eellison commented Mar 22, 2019 • edited

t-vi commented Mar 22, 2019

eellison commented Mar 22, 2019

t-vi commented Mar 22, 2019

eellison left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

eellison Mar 22, 2019 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

wanchaol left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

t-vi Mar 23, 2019 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

t-vi Mar 23, 2019 • edited

Choose a reason for hiding this comment

t-vi commented Mar 23, 2019

facebook-github-bot left a comment

Choose a reason for hiding this comment

t-vi commented Mar 23, 2019

apaszke left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

eellison commented Mar 22, 2019 •

edited

eellison Mar 22, 2019 •

edited

t-vi Mar 23, 2019 •

edited

t-vi Mar 23, 2019 •

edited