Some improvements to nonzero post guard_size_oblivious

[ezyang]

Stack from ghstack (oldest at bottom):

• -> Some improvements to nonzero post guard_size_oblivious #122156

Prompted by #121571

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

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 1111339 with merge base 852111e ():

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

• pull / linux-focal-py3_8-clang9-xla / test (xla, 1, 1, linux.12xlarge, unstable) (gh)
  Build completed, 1 test FAILED, 19959 total actions

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

[ezyang]

@pytorchbot rebase

[pytorchmergebot]

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

[pytorchmergebot]

Successfully rebased gh/ezyang/2616/orig onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via ghstack checkout https://github.com/pytorch/pytorch/pull/122156)

[BowenBao]

This diff probably helps unblock you from ONNX CI failures.

diff --git a/test/onnx/test_fx_op_consistency.py b/test/onnx/test_fx_op_consistency.py
index dc90977101b..506eb25f9f2 100644
--- a/test/onnx/test_fx_op_consistency.py
+++ b/test/onnx/test_fx_op_consistency.py
@@ -1563,11 +1563,10 @@ SKIP_XFAIL_SUBTESTS: tuple[onnx_test_common.DecorateMeta, ...] = (
     ),
     xfail(
         "nonzero",
-        model_type=pytorch_test_common.TorchModelType.TORCH_EXPORT_EXPORTEDPROGRAM,
-        reason=onnx_test_common.reason_onnx_script_does_not_support(
-            "aten::_assert_async.msg",
-            "https://github.com/pytorch/pytorch/issues/112443",
-        ),
+        matcher=lambda sample: len(sample.input.shape) == 0
+        and sample.kwargs.get("as_tuple", False) is False,
+        reason="Output 'shape' do not match: torch.Size([0, 1]) != torch.Size([0, 0]).",
+        model_type=pytorch_test_common.TorchModelType.TORCH_NN_MODULE,
     ),
     xfail(
         "scatter_add",

[ezyang]

@pytorchbot merge

[pytorchmergebot]

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