Python 3.10 Union operator | support for JIT #109293

randolf-scholz · 2023-09-14T13:13:03Z

Duplicated the tests from test/jit/test_union.py into test/jit/test_union_pep604.py, using PEP604 style Unions
Exchanged custom get_args and get_origin with typing.get_args and typing.get_origin which have the same functionality and became part of the standard library in 3.8
Added utility function pep604union_to_union in tree_views.h which converts a BinOP("|") node into the corresponding Union. This function intercepts ScriptTypeParser::parseTypeFromExpr and ScriptTypeParser::parseTypeFromExprImpl and patches the expression.
There is a single failing test, I commented it out for the moment to see if CI complains about anything else. I tried several hours to figure out how to patch it, but I am not experienced with C++ development and debugging.

From what I could gather, the following fails:

    def test_union_optional_of_union_return(self):
        @torch.jit.script
        def fn() -> None | str | int:
            y: Optional[int | str] = "foo"
            return y

In the section:

pytorch/torch/csrc/jit/frontend/script_type_parser.cpp

Lines 232 to 243 in 75b954b

    
           TypePtr ScriptTypeParser::parseTypeFromExpr(const Expr& expr) const { 
        
             // the resolver needs to recursively resolve the expression, so to avoid 
        
             // resolving all type expr subtrees we only use it for the top level 
        
             // expression and base type names. 
        
             if (resolver_) { 
        
               if (auto typePtr = 
        
                       resolver_->resolveType(expr.range().text().str(), expr.range())) { 
        
                 return typePtr; 
        
               } 
        
             } 
        
             return parseTypeFromExprImpl(expr); 
        
           }

When using regular Union, the resolver path is taken, whereas with the patch pep604 union, resolveType doesn't work.

cc @ezyang @malfet @rgommers @xuzhao9 @gramster

pytorch-bot · 2023-09-14T13:13:05Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/109293

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (1 Unrelated Failure)

As of commit 5a8c408 with merge base 5ad1baf ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

linux-focal-cuda11.8-py3.10-gcc9 / test (distributed, 1, 3, linux.8xlarge.nvidia.gpu) (gh)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

randolf-scholz · 2023-09-20T09:22:48Z

@pytorchbot label "module: typing"

ezyang

okey dokey

ezyang · 2023-09-21T15:49:10Z

@pytorchbot merge

pytorchmergebot · 2023-09-21T15:50:50Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

randolf-scholz · 2023-09-21T16:33:55Z

@ezyang This shouldn't be merged yet, first the failure of the Optional simplification should be clarified (as per review request). (whether to resolve it or to simply don't care about it)

EDIT: I forget to make the review visible...

test/jit/test_union_pep604.py

randolf-scholz · 2023-09-21T16:42:55Z

@ezyang I couldn't quite figure out how to solve it. Either we find a solution or one can decide to not care about this error. Also if you haven't it would be good to check the C++ code as I am a novice with respect to C++ development.

ezyang · 2023-09-24T02:54:04Z

torch/csrc/jit/frontend/tree_views.h

+      expr.range(),
+      Var::create(expr.range(), Ident::create(expr.range(), "Union")),
+      List<Expr>::create(expr.range(), members));
+  return std::move(synthesised_union);


nit: just return Subscript::create directly. There's a long thing about return on std::move impeding copy elision but iirc Exprs are cheap to move so it doesn't really matter.

I added this after it was suggested in failed CI: https://github.com/pytorch/pytorch/actions/runs/6209430790/job/16857935849#step:10:13787

ezyang · 2023-09-24T02:57:13Z

I'm willing to accept this with bugs. It will just mean some valid programs fail to typecheck, right? You can always figure it out later and send us a patch.

We don't have dedicated TS cycles so I am mostly estimating that this feature is relatively low risk to add. I also am not sure how to solve your particular problem, but if this would be useful to you it's better to just get this in.

randolf-scholz · 2023-09-24T08:30:58Z

I'm willing to accept this with bugs. It will just mean some valid programs fail to typecheck, right? You can always figure it out later and send us a patch.

We don't have dedicated TS cycles so I am mostly estimating that this feature is relatively low risk to add. I also am not sure how to solve your particular problem, but if this would be useful to you it's better to just get this in.

It means some valid programs fail JIT-compilation, specifically, this test fails (in the duplicated tests from test/jit/test_union.py):

    def test_union_optional_of_union_is_flattened(self):
        @torch.jit.script
        def fn(flag: int) -> str | int | None:
            y: int | str | None = "foo"
            if flag == 0:
                x: Optional[int | str] = y
            elif flag == 1:
                x: Optional[int | str] = 1
            else:
                x: Optional[int | str] = None
            return x

with the following error message upon attempting to jit.script the function:

Return value was annotated as having type Union[int, NoneType, str]
but is actually of type Optional[Union[int, str]]:

It succeeds if |-unions are replaced by typing.Union, or if the Optional is replaced by None | ....

ezyang · 2023-09-25T02:42:42Z

Yeah, this is OK by me. Essentially, you are only strictly expanding the set of programs we accept, there are some programs we don't manage to get but they are easy to work around.

ezyang · 2023-09-25T15:28:48Z

@pytorchbot merge

pytorchmergebot · 2023-09-25T15:30:59Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

pytorchmergebot · 2023-09-25T15:31:05Z

Merge failed

Reason: 1 mandatory check(s) failed. The first few are:

pull / linux-focal-cuda11.8-py3.10-gcc9 / test (distributed, 1, 3, linux.8xlarge.nvidia.gpu)

Dig deeper by viewing the failures on hud

Details for Dev Infra team

Raised by workflow job

Failing merge rule: Core Maintainers

ezyang · 2023-09-25T15:33:23Z

@pytorchbot merge -f "spurious failure"

pytorchmergebot · 2023-09-25T15:35:49Z

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use -f as last resort and instead consider -i/--ignore-current to continue the merge ignoring current failures. This will allow currently pending tests to finish and report signal before the merge.

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

wip JIT PEP604

7c2d2b4

pytorch-bot bot added the release notes: jit release notes category label Sep 14, 2023

pytorchbot added the open source label Sep 14, 2023

randolf-scholz added 10 commits September 14, 2023 22:53

WIP

a011dc5

wip

e7ceb2f

wip

39fe05c

Merge branch 'main' into jit_script_union

5226a5c

wip

477c34e

added utility function to convert pep604 union to classical union

4ce7690

fixed some linter issues

2ed4bc0

Merge branch 'pytorch:main' into jit_script_union

b781e2b

Merge branch 'main' into jit_script_union

a381e86

added std::move

f769047

randolf-scholz marked this pull request as ready for review September 17, 2023 19:02

randolf-scholz changed the title ~~wip JIT PEP604~~ Python 3.10 Union operator | support for JIT Sep 17, 2023

pytorch-bot bot added the module: typing Related to mypy type annotations label Sep 20, 2023

soulitzer requested a review from davidberard98 September 20, 2023 15:32

soulitzer added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Sep 20, 2023

davidberard98 requested a review from gmagogsfm September 20, 2023 17:16

ezyang approved these changes Sep 21, 2023

View reviewed changes

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Sep 21, 2023

pytorchmergebot added the merging label Sep 21, 2023

randolf-scholz commented Sep 21, 2023

View reviewed changes

test/jit/test_union_pep604.py Outdated Show resolved Hide resolved

randolf-scholz closed this Sep 21, 2023

randolf-scholz reopened this Sep 21, 2023

This comment was marked as spam.

Sign in to view

randolf-scholz closed this Sep 23, 2023

randolf-scholz reopened this Sep 23, 2023

randolf-scholz added 2 commits September 23, 2023 14:12

Merge branch 'main' into jit_script_union

b6607e7

Merge branch 'main' into jit_script_union

0e05d24

ezyang reviewed Sep 24, 2023

View reviewed changes

ezyang requested a review from eellison September 24, 2023 02:56

Merge branch 'main' into jit_script_union

01dcbc1

marked tests with expectedFailure instead of skip

5a8c408

pytorchmergebot added the merging label Sep 25, 2023

pytorchmergebot removed the merging label Sep 25, 2023

ezyang added the topic: new features topic category label Sep 25, 2023

pytorchmergebot added the merging label Sep 25, 2023

pytorchmergebot added Merged and removed merging labels Sep 25, 2023

pytorchmergebot closed this in 837272f Sep 25, 2023

tringwald mentioned this pull request Nov 30, 2023

Torchscript should support union type by "|" introduced in python3.10. #114755

Closed

laserkelvin mentioned this pull request Feb 13, 2024

[Bug]: MACE fails to JIT compile with Python 3.10 union type hints IntelLabs/matsciml#123

Closed

akihironitta mentioned this pull request Apr 18, 2025

Type hint error for python 3.12.2(?) pyg-team/pytorch_geometric#10138

Open

	TypePtr ScriptTypeParser::parseTypeFromExpr(const Expr& expr) const {
	// the resolver needs to recursively resolve the expression, so to avoid
	// resolving all type expr subtrees we only use it for the top level
	// expression and base type names.
	if (resolver_) {
	if (auto typePtr =
	resolver_->resolveType(expr.range().text().str(), expr.range())) {
	return typePtr;
	}
	}
	return parseTypeFromExprImpl(expr);
	}

Python 3.10 Union operator | support for JIT #109293

Python 3.10 Union operator | support for JIT #109293

Uh oh!

Conversation

randolf-scholz commented Sep 14, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Sep 14, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/109293

✅ You can merge normally! (1 Unrelated Failure)

Uh oh!

randolf-scholz commented Sep 20, 2023

Uh oh!

ezyang left a comment

Choose a reason for hiding this comment

Uh oh!

ezyang commented Sep 21, 2023

Uh oh!

pytorchmergebot commented Sep 21, 2023

Merge started

Uh oh!

randolf-scholz commented Sep 21, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

randolf-scholz commented Sep 21, 2023

Uh oh!

This comment was marked as spam.

Uh oh!

ezyang Sep 24, 2023

Choose a reason for hiding this comment

Uh oh!

randolf-scholz Sep 24, 2023

Choose a reason for hiding this comment

Uh oh!

ezyang commented Sep 24, 2023

Uh oh!

randolf-scholz commented Sep 24, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ezyang commented Sep 25, 2023

Uh oh!

ezyang commented Sep 25, 2023

Uh oh!

pytorchmergebot commented Sep 25, 2023

Merge started

Uh oh!

pytorchmergebot commented Sep 25, 2023

Merge failed

Uh oh!

ezyang commented Sep 25, 2023

Uh oh!

pytorchmergebot commented Sep 25, 2023

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

randolf-scholz commented Sep 14, 2023 •

edited

Loading

pytorch-bot bot commented Sep 14, 2023 •

edited

Loading

randolf-scholz commented Sep 21, 2023 •

edited

Loading

randolf-scholz commented Sep 24, 2023 •

edited

Loading