fix aminmax output resize issue when input is a zero dimension tensor

[mingfeima]

Stack from ghstack (oldest at bottom):

• -> fix aminmax output resize issue when input is a zero dimension tensor #96171

Fix #96042

before

>>> torch.aminmax(torch.tensor(1, device='cpu'), dim=0, keepdim=True)
__main__:1: UserWarning: An output with one or more elements was resized since it had shape [], which does not match the required output shape [1]. This behavior is deprecated, and in a future PyTorch release outputs will not be resized unless they have zero elements. You can explicitly reuse an out tensor t by resizing it, inplace, to zero elements with t.resize_(0). (Triggered internally at ../aten/src/ATen/native/Resize.cpp:24.)
torch.return_types.aminmax(
min=tensor([1]),
max=tensor([1]))
>>> torch.aminmax(torch.tensor(1, device='cpu'), dim=0, keepdim=False)
torch.return_types.aminmax(
min=tensor(1),
max=tensor(1))

after

>>> torch.aminmax(torch.tensor(1, device='cpu'), dim=0, keepdim=True)
torch.return_types.aminmax(
min=tensor(1),
max=tensor(1))
>>> torch.aminmax(torch.tensor(1, device='cpu'), dim=0, keepdim=False)
torch.return_types.aminmax(
min=tensor(1),
max=tensor(1))

Marked the following test as expected_fail:
test_vmap.py TestVmapOperatorsOpInfoCPU.test_op_has_batch_rule_aminmax_cpu_float32

Given input shape of (2), the loop out is shape (2), the batched vmap out is (2, 1), which mismatched.
The loop out will calculate twice on a tensor shape of ( ): without this patch, the output is (1), and then stacked into (2, 1); with this patch, the output is ( ), then stacked into (2).

cc @jgong5 @XiaobingSuper @sanchitintel @ashokei @jingxu10

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 Failures

As of commit efc5fa6:

NEW FAILURES - The following jobs have failed:

• win-vs2019-cpu-py3 / test (default, 2, 2, windows.4xlarge) (gh)

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

[ngimel]

You'd also need to remove special handling in decomp now that CPU doesn't have inconsistent behavior

[ngimel]

cc @zou3519 for vmap errors. If 0d tensor is vmapped, what's the expected result? Here it seems vmap produces (2,1) whereas loop correctly produces (2).
@mingfeima I think we can skip failing tests for now.

[zou3519]

cc @zou3519 for vmap errors. If 0d tensor is vmapped, what's the expected result? Here it seems vmap produces (2,1) whereas loop correctly produces (2).

The expected result is whatever the loop produces (so it sounds like 2).

From the comment over here, it sounds like we can just add a REDUCTION_WITH_KEEPDIM_ARG(aminmax) like this and delete this.

This is not too much of a change, if you're interested in handing in this PR @mingfeima. Otherwise I can handle it later after this PR lands

[mingfeima]

cc @zou3519 for vmap errors. If 0d tensor is vmapped, what's the expected result? Here it seems vmap produces (2,1) whereas loop correctly produces (2).

The expected result is whatever the loop produces (so it sounds like 2).

From the comment over here, it sounds like we can just add a REDUCTION_WITH_KEEPDIM_ARG(aminmax) like this and delete this.

This is not too much of a change, if you're interested in handing in this PR @mingfeima. Otherwise I can handle it later after this PR lands

@zou3519 Thanks for the input! Update this patch according to your proposal.

[zou3519]

Thank you!

[ngimel]

@pytorchbot merge -f "jit failure unrelated"

[pytorchmergebot]

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes).

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