Use expect_true to make split with unbacked sizes work. #106788

ezyang · 2023-08-08T14:37:45Z

Stack from ghstack (oldest at bottom):

This pattern shows up in torchrec KeyedJaggedTensor. Most
of the change in this PR is mechanical: whenever we failed
an unbacked symint test due to just error checking, replace the
conditional with something that calls expect_true (e.g.,
torch._check or TORCH_SYM_CHECK).

Some of the changes are a bit more nuanced, I've commented on the PR
accordingly.

Signed-off-by: Edward Z. Yang ezyang@meta.com

cc @voznesenskym @penguinwu @anijain2305 @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @ipiszy @chenyang78 @aakhundov

This pattern shows up in torchrec KeyedJaggedTensor. Most of the change in this PR is mechanical: whenever we failed an unbacked symint test due to just error checking, replace the conditional with something that calls expect_true (e.g., torch._check or TORCH_SYM_CHECK). Some of the changes are a bit more nuanced, I've commented on the PR accordingly. Signed-off-by: Edward Z. Yang <ezyang@meta.com> [ghstack-poisoned]

pytorch-bot · 2023-08-08T14:37:53Z

🔗 Helpful Links

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

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

✅ 1 Unrelated Failure

As of commit 0526cb2:

UNSTABLE - The following job failed but was likely due to flakiness present on trunk and has been marked as unstable:

cuda12.1-py3.10-gcc9-sm86 / test (inductor, 1, 1, linux.g5.4xlarge.nvidia.gpu, unstable) (gh)

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

This pattern shows up in torchrec KeyedJaggedTensor. Most of the change in this PR is mechanical: whenever we failed an unbacked symint test due to just error checking, replace the conditional with something that calls expect_true (e.g., torch._check or TORCH_SYM_CHECK). Some of the changes are a bit more nuanced, I've commented on the PR accordingly. Signed-off-by: Edward Z. Yang <ezyangmeta.com> ghstack-source-id: dcdfbb704ed64b9a95b02157d41fb55bd383c1f4 Pull Request resolved: #106788

ezyang · 2023-08-08T14:39:21Z

aten/src/ATen/native/TensorShape.cpp

-    start = maybe_wrap_dim(start, cur_size);
+  TORCH_SYM_CHECK((-cur_size).sym_le(start).sym_and(start.sym_le(cur_size)), "narrow(): start must be within [-cur_size, cur_size]");
+  if (start < 0) {
+    start = start + cur_size;


The start != cur_size test cannot be sym'ified since one branch is not an error condition. What I did here was inlined maybe_wrap_dim here, which conventionally tests -cur_size <= start < cur_size; so you can make it work without extra branching by just changing this condition to -cur_size <= start <= cur_size.

I can split out the inlining into its own PR if people would prefer.

ezyang · 2023-08-08T14:40:29Z

torch/__init__.py

@@ -948,7 +948,7 @@ def _check_with(error_type, cond: Union[builtins.bool, SymBool], message: Callab
    if not isinstance(cond, (builtins.bool, torch.SymBool)):
        raise TypeError(f'cond must be a bool, but got {type(cond)}')

-    if cond:
+    if torch.fx.experimental.symbolic_shapes.expect_true(cond):


This turns on expect_true for all our decomps/meta functions, pretty nice. (I couldn't do the same trick in C++; TORCH_CHECK is a very delicate macro and it was hard to insert expect_true without breaking some sites, and there's also the problem that operator< and friends actually return bool not SymBool).

nb. You can do a horrible, horrible thing, and locally overwrite the meaning of < inside TORCH_CHECK via a macro so that you automagically go from using operator< to lt.

Also, this line is fantastic. Gotta love having just one entrypoint for a given thing.

ezyang · 2023-08-08T14:41:32Z

torch/_decomp/decompositions.py

        start_val = sizes[dim]

    if end_val < start_val:
        end_val = start_val
-    elif end_val >= sizes[dim]:
+    elif end_val > sizes[dim]:


These two changes don't actually change semantics, but they're pretty important: frequently, we will know that end_val is <= sizes[dim], but we don't know if it == sizes[dim] or not. By branching only if it is truly out of bounds, we can statically determine which branch we go down. This is enough for split.

This pattern shows up in torchrec KeyedJaggedTensor. Most of the change in this PR is mechanical: whenever we failed an unbacked symint test due to just error checking, replace the conditional with something that calls expect_true (e.g., torch._check or TORCH_SYM_CHECK). Some of the changes are a bit more nuanced, I've commented on the PR accordingly. Signed-off-by: Edward Z. Yang <ezyangmeta.com> [ghstack-poisoned]

This pattern shows up in torchrec KeyedJaggedTensor. Most of the change in this PR is mechanical: whenever we failed an unbacked symint test due to just error checking, replace the conditional with something that calls expect_true (e.g., torch._check or TORCH_SYM_CHECK). Some of the changes are a bit more nuanced, I've commented on the PR accordingly. Signed-off-by: Edward Z. Yang <ezyangmeta.com> ghstack-source-id: c93177d89e1a1d6c15f2b0765c77c9f6282ec891 Pull Request resolved: #106788

This pattern shows up in torchrec KeyedJaggedTensor. Most of the change in this PR is mechanical: whenever we failed an unbacked symint test due to just error checking, replace the conditional with something that calls expect_true (e.g., torch._check or TORCH_SYM_CHECK). Some of the changes are a bit more nuanced, I've commented on the PR accordingly. Signed-off-by: Edward Z. Yang <ezyangmeta.com> [ghstack-poisoned]

This pattern shows up in torchrec KeyedJaggedTensor. Most of the change in this PR is mechanical: whenever we failed an unbacked symint test due to just error checking, replace the conditional with something that calls expect_true (e.g., torch._check or TORCH_SYM_CHECK). Some of the changes are a bit more nuanced, I've commented on the PR accordingly. Signed-off-by: Edward Z. Yang <ezyangmeta.com> ghstack-source-id: ee86135ddd20bcf671682f5c5bdb728421dc701e Pull Request resolved: #106788

This pattern shows up in torchrec KeyedJaggedTensor. Most of the change in this PR is mechanical: whenever we failed an unbacked symint test due to just error checking, replace the conditional with something that calls expect_true (e.g., torch._check or TORCH_SYM_CHECK). Some of the changes are a bit more nuanced, I've commented on the PR accordingly. Signed-off-by: Edward Z. Yang <ezyangmeta.com> [ghstack-poisoned]

This pattern shows up in torchrec KeyedJaggedTensor. Most of the change in this PR is mechanical: whenever we failed an unbacked symint test due to just error checking, replace the conditional with something that calls expect_true (e.g., torch._check or TORCH_SYM_CHECK). Some of the changes are a bit more nuanced, I've commented on the PR accordingly. Signed-off-by: Edward Z. Yang <ezyangmeta.com> ghstack-source-id: 11d5692b882a2825312335b046cf1c6d8e798e95 Pull Request resolved: #106788

torch/_refs/__init__.py

This pattern shows up in torchrec KeyedJaggedTensor. Most of the change in this PR is mechanical: whenever we failed an unbacked symint test due to just error checking, replace the conditional with something that calls expect_true (e.g., torch._check or TORCH_SYM_CHECK). Some of the changes are a bit more nuanced, I've commented on the PR accordingly. Signed-off-by: Edward Z. Yang <ezyangmeta.com> cc voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx ipiszy chenyang78 aakhundov [ghstack-poisoned]

This pattern shows up in torchrec KeyedJaggedTensor. Most of the change in this PR is mechanical: whenever we failed an unbacked symint test due to just error checking, replace the conditional with something that calls expect_true (e.g., torch._check or TORCH_SYM_CHECK). Some of the changes are a bit more nuanced, I've commented on the PR accordingly. Signed-off-by: Edward Z. Yang <ezyangmeta.com> ghstack-source-id: ff5b6867ba795ec6ba84d921fa4caf98f0325bc8 Pull Request resolved: #106788

This pattern shows up in torchrec KeyedJaggedTensor. Most of the change in this PR is mechanical: whenever we failed an unbacked symint test due to just error checking, replace the conditional with something that calls expect_true (e.g., torch._check or TORCH_SYM_CHECK). Some of the changes are a bit more nuanced, I've commented on the PR accordingly. Signed-off-by: Edward Z. Yang <ezyangmeta.com> cc voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx ipiszy chenyang78 aakhundov [ghstack-poisoned]

This pattern shows up in torchrec KeyedJaggedTensor. Most of the change in this PR is mechanical: whenever we failed an unbacked symint test due to just error checking, replace the conditional with something that calls expect_true (e.g., torch._check or TORCH_SYM_CHECK). Some of the changes are a bit more nuanced, I've commented on the PR accordingly. Signed-off-by: Edward Z. Yang <ezyangmeta.com> ghstack-source-id: 41bea703b17b260d7bee0d19359d9866f5c2bde5 Pull Request resolved: #106788

aten/src/ATen/native/TensorShape.cpp

This pattern shows up in torchrec KeyedJaggedTensor. Most of the change in this PR is mechanical: whenever we failed an unbacked symint test due to just error checking, replace the conditional with something that calls expect_true (e.g., torch._check or TORCH_SYM_CHECK). Some of the changes are a bit more nuanced, I've commented on the PR accordingly. Signed-off-by: Edward Z. Yang <ezyangmeta.com> cc voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx ipiszy chenyang78 aakhundov [ghstack-poisoned]

This pattern shows up in torchrec KeyedJaggedTensor. Most of the change in this PR is mechanical: whenever we failed an unbacked symint test due to just error checking, replace the conditional with something that calls expect_true (e.g., torch._check or TORCH_SYM_CHECK). Some of the changes are a bit more nuanced, I've commented on the PR accordingly. Signed-off-by: Edward Z. Yang <ezyangmeta.com> ghstack-source-id: bf62bf051e6a82a35c44830eb8d1b9aa98fe9e99 Pull Request resolved: #106788

This pattern shows up in torchrec KeyedJaggedTensor. Most of the change in this PR is mechanical: whenever we failed an unbacked symint test due to just error checking, replace the conditional with something that calls expect_true (e.g., torch._check or TORCH_SYM_CHECK). Some of the changes are a bit more nuanced, I've commented on the PR accordingly. Signed-off-by: Edward Z. Yang <ezyangmeta.com> cc voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx ipiszy chenyang78 aakhundov [ghstack-poisoned]

This pattern shows up in torchrec KeyedJaggedTensor. Most of the change in this PR is mechanical: whenever we failed an unbacked symint test due to just error checking, replace the conditional with something that calls expect_true (e.g., torch._check or TORCH_SYM_CHECK). Some of the changes are a bit more nuanced, I've commented on the PR accordingly. Signed-off-by: Edward Z. Yang <ezyangmeta.com> ghstack-source-id: b1ba0de493b3141f54a844f683cdf94620d721ab Pull Request resolved: #106788

ezyang · 2023-08-15T20:27:13Z

@pytorchbot merge -f "unstable only"

pytorchmergebot · 2023-08-15T20:31:24Z

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

This pattern shows up in torchrec KeyedJaggedTensor. Most of the change in this PR is mechanical: whenever we failed an unbacked symint test due to just error checking, replace the conditional with something that calls expect_true (e.g., torch._check or TORCH_SYM_CHECK). Some of the changes are a bit more nuanced, I've commented on the PR accordingly. Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: pytorch#106788 Approved by: https://github.com/lezcano ghstack dependencies: pytorch#106720

See #106788 for context. I think I don't actually need this for anything real, but this is pretty mild so might as well. Signed-off-by: Edward Z. Yang <ezyangmeta.com> ghstack-source-id: b905ba0dbcce50af5a2411cedc89c051add745ea Pull Request resolved: #107785

See #106788 for context. I think I don't actually need this for anything real, but this is pretty mild so might as well. Signed-off-by: Edward Z. Yang <ezyang@meta.com> [ghstack-poisoned]

…7785) See #106788 for context. I think I don't actually need this for anything real, but this is pretty mild so might as well. Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: #107785 Approved by: https://github.com/albanD

ezyang mentioned this pull request Aug 8, 2023

Add expect_true for irrefutable guards #106720

Closed

github-actions bot requested a review from albanD August 8, 2023 14:38

github-actions bot added the ciflow/inductor label Aug 8, 2023

github-actions bot requested review from antoniojkim, bdhirsh, jbschlosser, miladm, SherlockNoMad, voznesenskym and wconstab August 8, 2023 14:38

ezyang requested a review from tugsbayasgalan August 8, 2023 14:38

ezyang commented Aug 8, 2023

View reviewed changes

albanD removed their request for review August 8, 2023 14:48

ezyang requested a review from ysiraichi August 8, 2023 15:00

ezyang requested review from mruberry and ngimel as code owners August 8, 2023 16:38

ezyang mentioned this pull request Aug 8, 2023

Use expect tests for error inputs #106828

Open

ezyang added ciflow/trunk Trigger trunk jobs on your pull request topic: not user facing topic category labels Aug 9, 2023

github-actions bot added the module: dynamo label Aug 9, 2023

BowenBao reviewed Aug 11, 2023

View reviewed changes

torch/_refs/__init__.py Show resolved Hide resolved

ezyang mentioned this pull request Aug 14, 2023

[WIP] Force KJT list inputs to be unbacked ints #107156

Closed

lezcano reviewed Aug 14, 2023

View reviewed changes

aten/src/ATen/native/TensorShape.cpp Show resolved Hide resolved

lezcano approved these changes Aug 15, 2023

View reviewed changes

pytorchmergebot added the merging label Aug 15, 2023

pytorchmergebot added Merged and removed merging labels Aug 15, 2023

pytorchmergebot closed this in 5673c08 Aug 15, 2023

facebook-github-bot deleted the gh/ezyang/2289/head branch August 19, 2023 14:16

ezyang mentioned this pull request Aug 23, 2023

Use TORCH_SYM_CHECK for check_size_nonnegative on SymIntArrayRef #107785

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Use expect_true to make split with unbacked sizes work. #106788

Use expect_true to make split with unbacked sizes work. #106788

ezyang commented Aug 8, 2023 •

edited

pytorch-bot bot commented Aug 8, 2023 •

edited

ezyang Aug 8, 2023

ezyang Aug 8, 2023

ezyang Aug 8, 2023

lezcano Aug 15, 2023 •

edited

lezcano Aug 15, 2023

ezyang Aug 8, 2023

ezyang commented Aug 15, 2023

pytorchmergebot commented Aug 15, 2023

Use expect_true to make split with unbacked sizes work. #106788

Use expect_true to make split with unbacked sizes work. #106788

Conversation

ezyang commented Aug 8, 2023 • edited

pytorch-bot bot commented Aug 8, 2023 • edited

🔗 Helpful Links

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

✅ 1 Unrelated Failure

ezyang Aug 8, 2023

Choose a reason for hiding this comment

ezyang Aug 8, 2023

Choose a reason for hiding this comment

ezyang Aug 8, 2023

Choose a reason for hiding this comment

lezcano Aug 15, 2023 • edited

Choose a reason for hiding this comment

lezcano Aug 15, 2023

Choose a reason for hiding this comment

ezyang Aug 8, 2023

Choose a reason for hiding this comment

ezyang commented Aug 15, 2023

pytorchmergebot commented Aug 15, 2023

Merge started

ezyang commented Aug 8, 2023 •

edited

pytorch-bot bot commented Aug 8, 2023 •

edited

lezcano Aug 15, 2023 •

edited